From 171a475db8cd8a43f94382a5b302d1d8015effc8 Mon Sep 17 00:00:00 2001 From: RenQin <45731309+soonera@users.noreply.github.com> Date: Tue, 2 Aug 2022 16:25:09 +0800 Subject: [PATCH] [Feature]: save mae visualization results (#388) --- docs/en/get_started.md | 10 ++++++---- docs/zh_cn/get_started.md | 32 +++++++++++++++++--------------- tools/misc/mae_visualization.py | 13 +++++++------ 3 files changed, 30 insertions(+), 25 deletions(-) diff --git a/docs/en/get_started.md b/docs/en/get_started.md index efc204518..05f280565 100644 --- a/docs/en/get_started.md +++ b/docs/en/get_started.md @@ -11,6 +11,7 @@ - [Count number of parameters](#count-number-of-parameters) - [Publish a model](#publish-a-model) - [Use t-SNE](#use-t-sne) + - [MAE Visualization](#mae-visualization) - [Reproducibility](#reproducibility) This page provides basic tutorials about the usage of MMSelfSup. For installation instructions, please see [install.md](install.md). @@ -171,20 +172,21 @@ Arguments: We provide a tool to visualize the mask and reconstruction image of MAE model. ```shell -python tools/misc/mae_visualization.py ${IMG} ${CONFIG_FILE} ${CKPT_PATH} --device ${DEVICE} +python tools/misc/mae_visualization.py ${IMG_PATH} ${CONFIG_FILE} ${CKPT_PATH} ${OUT_FILE} --device ${DEVICE} ``` -参数: +Arguments: -- `IMG`: an image path used for visualization. +- `IMG_PATH`: an image path used for visualization. - `CONFIG_FILE`: config file for the pre-trained model. - `CKPT_PATH`: the path of model's checkpoint. +- `OUT_FILE`: the image path used for visualization results. - `DEVICE`: device used for inference. An example: ```shell -python tools/misc/mae_visualization.py tests/data/color.jpg configs/selfsup/mae/mae_vit-base-p16_8xb512-coslr-400e_in1k.py mae_epoch_400.pth --device 'cuda:0' +python tools/misc/mae_visualization.py tests/data/color.jpg configs/selfsup/mae/mae_vit-base-p16_8xb512-coslr-400e_in1k.py mae_epoch_400.pth results.jpg --device 'cuda:0' ``` ### Reproducibility diff --git a/docs/zh_cn/get_started.md b/docs/zh_cn/get_started.md index 51525f6a6..7ce8bc99f 100644 --- a/docs/zh_cn/get_started.md +++ b/docs/zh_cn/get_started.md @@ -1,17 +1,18 @@ # 基础教程 -- [基础教程](#%E5%9F%BA%E7%A1%80%E6%95%99%E7%A8%8B) - - [训练已有的算法](#%E8%AE%AD%E7%BB%83%E5%B7%B2%E6%9C%89%E7%9A%84%E7%AE%97%E6%B3%95) - - [使用 CPU 训练](#%E4%BD%BF%E7%94%A8-cpu-%E8%AE%AD%E7%BB%83) - - [使用 单张/多张 显卡训练](#%E4%BD%BF%E7%94%A8-%E5%8D%95%E5%BC%A0%E5%A4%9A%E5%BC%A0-%E6%98%BE%E5%8D%A1%E8%AE%AD%E7%BB%83) - - [使用多台机器训练](#%E4%BD%BF%E7%94%A8%E5%A4%9A%E5%8F%B0%E6%9C%BA%E5%99%A8%E8%AE%AD%E7%BB%83) - - [在一台机器上启动多个任务](#%E5%9C%A8%E4%B8%80%E5%8F%B0%E6%9C%BA%E5%99%A8%E4%B8%8A%E5%90%AF%E5%8A%A8%E5%A4%9A%E4%B8%AA%E4%BB%BB%E5%8A%A1) - - [基准测试](#%E5%9F%BA%E5%87%86%E6%B5%8B%E8%AF%95) - - [工具和建议](#%E5%B7%A5%E5%85%B7%E5%92%8C%E5%BB%BA%E8%AE%AE) - - [统计模型的参数](#%E7%BB%9F%E8%AE%A1%E6%A8%A1%E5%9E%8B%E7%9A%84%E5%8F%82%E6%95%B0) - - [发布模型](#%E5%8F%91%E5%B8%83%E6%A8%A1%E5%9E%8B) - - [使用 t-SNE 来做模型可视化](#%E4%BD%BF%E7%94%A8-t-sne-%E6%9D%A5%E5%81%9A%E6%A8%A1%E5%9E%8B%E5%8F%AF%E8%A7%86%E5%8C%96) - - [可复现性](#%E5%8F%AF%E5%A4%8D%E7%8E%B0%E6%80%A7) +- [基础教程](#基础教程) + - [训练已有的算法](#训练已有的算法) + - [使用 CPU 训练](#使用-cpu-训练) + - [使用 单张/多张 显卡训练](#使用-单张多张-显卡训练) + - [使用多台机器训练](#使用多台机器训练) + - [在一台机器上启动多个任务](#在一台机器上启动多个任务) + - [基准测试](#基准测试) + - [工具和建议](#工具和建议) + - [统计模型的参数](#统计模型的参数) + - [发布模型](#发布模型) + - [使用 t-SNE 来做模型可视化](#使用-t-sne-来做模型可视化) + - [MAE 可视化](#mae-可视化) + - [可复现性](#可复现性) 本文档提供 MMSelfSup 相关用法的基础教程。 如果您对如何安装 MMSelfSup 以及其相关依赖库有疑问, 请参考 [安装文档](install.md). @@ -169,20 +170,21 @@ python tools/analysis_tools/visualize_tsne.py ${CONFIG_FILE} --checkpoint ${CKPT 我们提供了一个对 MAE 掩码效果和重建效果可视化可视化的方法: ```shell -python tools/misc/mae_visualization.py ${IMG} ${CONFIG_FILE} ${CKPT_PATH} --device ${DEVICE} +python tools/misc/mae_visualization.py ${IMG_PATH} ${CONFIG_FILE} ${CKPT_PATH} ${OUT_FILE} --device ${DEVICE} ``` 参数: -- `IMG`: 用于可视化的图片 +- `IMG_PATH`: 用于可视化的图片 - `CONFIG_FILE`: 训练预训练模型的参数配置文件. - `CKPT_PATH`: 预训练模型的路径. +- `OUT_FILE`: 用于保存可视化结果的图片路径 - `DEVICE`: 用于推理的设备. 示例: ```shell -python tools/misc/mae_visualization.py tests/data/color.jpg configs/selfsup/mae/mae_vit-base-p16_8xb512-coslr-400e_in1k.py mae_epoch_400.pth --device 'cuda:0' +python tools/misc/mae_visualization.py tests/data/color.jpg configs/selfsup/mae/mae_vit-base-p16_8xb512-coslr-400e_in1k.py mae_epoch_400.pth results.jpg --device 'cuda:0' ``` ### 可复现性 diff --git a/tools/misc/mae_visualization.py b/tools/misc/mae_visualization.py index ff38d2d53..22cd9ee9e 100644 --- a/tools/misc/mae_visualization.py +++ b/tools/misc/mae_visualization.py @@ -27,8 +27,8 @@ def show_image(image: torch.Tensor, title: str = '') -> None: return -def show_images(x: torch.Tensor, im_masked: torch.Tensor, y: torch.Tensor, - im_paste: torch.Tensor) -> None: +def save_images(x: torch.Tensor, im_masked: torch.Tensor, y: torch.Tensor, + im_paste: torch.Tensor, out_file: str) -> None: # make the plt figure larger plt.rcParams['figure.figsize'] = [24, 6] @@ -44,7 +44,7 @@ def show_images(x: torch.Tensor, im_masked: torch.Tensor, y: torch.Tensor, plt.subplot(1, 4, 4) show_image(im_paste, 'reconstruction + visible') - plt.show() + plt.savefig(out_file) def post_process( @@ -60,9 +60,10 @@ def post_process( def main(): parser = ArgumentParser() - parser.add_argument('img', help='Image file') + parser.add_argument('img_path', help='Image file path') parser.add_argument('config', help='MAE Config file') parser.add_argument('checkpoint', help='Checkpoint file') + parser.add_argument('out_file', help='The output image file path') parser.add_argument( '--device', default='cuda:0', help='Device used for inference') args = parser.parse_args() @@ -83,10 +84,10 @@ def main(): dict(type='Normalize', **img_norm_cfg), ])) - img = Image.open(args.img) + img = Image.open(args.img_path) img, (mask, pred) = inference_model(model, img) x, im_masked, y, im_paste = post_process(img, pred, mask) - show_images(x, im_masked, y, im_paste) + save_images(x, im_masked, y, im_paste, args.out_file) if __name__ == '__main__':