open-mmlab · fangyixiao18 · Aug 2, 2022 · Aug 1, 2022
diff --git 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
@@ -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
@@ -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__':