[Enhance] Refactor imshow_det_rbboxes (#187)

* refactor show result * add test_palette * add api visualization * 使用 Colaboratory 创建 * fix lint * 使用 Colaboratory 创建 * fix lint * fix name add palette to huge demo * add sar, hrsc, hrsc_classwise PALETTE * fix typo
open-mmlab · Apr 13, 2022 · 3a2dfbf · 3a2dfbf
1 parent 28ee073
commit 3a2dfbf
Show file tree

Hide file tree

Showing 14 changed files with 559 additions and 291 deletions.
diff --git a/demo/MMRotate_Tutorial.ipynb b/demo/MMRotate_Tutorial.ipynb
diff --git a/demo/README.md b/demo/README.md
@@ -6,8 +6,7 @@ We provide a demo script to test a single image.
 python demo/image_demo.py \
     ${IMG_ROOT} \
     ${CONFIG_FILE} \
-    ${CHECKPOINT_FILE} \
-    ${OUTPUT_ROOT}]
+    ${CHECKPOINT_FILE}
 ```
 
 Examples:
@@ -16,6 +15,5 @@ Examples:
     python demo/image_demo.py \
         demo/demo.jpg \
         work_dirs/oriented_rcnn_r50_fpn_1x_dota_v3/oriented_rcnn_r50_fpn_1x_dota_v3.py \
-        work_dirs/oriented_rcnn_r50_fpn_1x_dota_v3/epoch_12.pth \
-        demo/vis.jpg
+        work_dirs/oriented_rcnn_r50_fpn_1x_dota_v3/epoch_12.pth
 ```
diff --git a/demo/huge_image_demo.py b/demo/huge_image_demo.py
@@ -48,6 +48,11 @@ def parse_args():
         help='IoU threshould for merging results')
     parser.add_argument(
         '--device', default='cuda:0', help='Device used for inference')
+    parser.add_argument(
+        '--palette',
+        default='dota',
+        choices=['dota', 'sar', 'hrsc', 'hrsc_classwise', 'random'],
+        help='Color palette used for visualization')
     parser.add_argument(
         '--score-thr', type=float, default=0.3, help='bbox score threshold')
     args = parser.parse_args()
@@ -62,7 +67,12 @@ def main(args):
                                            args.patch_steps, args.img_ratios,
                                            args.merge_iou_thr)
     # show the results
-    show_result_pyplot(model, args.img, result, score_thr=args.score_thr)
+    show_result_pyplot(
+        model,
+        args.img,
+        result,
+        palette=args.palette,
+        score_thr=args.score_thr)
 
 
 if __name__ == '__main__':

diff --git a/demo/image_demo.py b/demo/image_demo.py
@@ -9,39 +9,49 @@
 python demo/image_demo.py \
     demo/demo.jpg \
     configs/oriented_rcnn/oriented_rcnn_r50_fpn_1x_dota_le90.py \
-    checkpoint/oriented_rcnn_r50_fpn_1x_dota_le90-6d2b2ce0.pth \
-    demo/vis.jpg
+    work_dirs/oriented_rcnn_r50_fpn_1x_dota_v3/epoch_12.pth
 ```
 """  # nowq
 
 from argparse import ArgumentParser
 
-from mmdet.apis import inference_detector, init_detector
+from mmdet.apis import inference_detector, init_detector, show_result_pyplot
 
 import mmrotate  # noqa: F401
 
 
-def main():
-    """Test a single image."""
+def parse_args():
     parser = ArgumentParser()
     parser.add_argument('img', help='Image file')
     parser.add_argument('config', help='Config file')
     parser.add_argument('checkpoint', help='Checkpoint file')
-    parser.add_argument('output', help='Output image')
     parser.add_argument(
         '--device', default='cuda:0', help='Device used for inference')
+    parser.add_argument(
+        '--palette',
+        default='dota',
+        choices=['dota', 'sar', 'hrsc', 'hrsc_classwise', 'random'],
+        help='Color palette used for visualization')
     parser.add_argument(
         '--score-thr', type=float, default=0.3, help='bbox score threshold')
     args = parser.parse_args()
+    return args
+
 
+def main(args):
     # build the model from a config file and a checkpoint file
     model = init_detector(args.config, args.checkpoint, device=args.device)
     # test a single image
     result = inference_detector(model, args.img)
     # show the results
-    model.show_result(
-        args.img, result, score_thr=args.score_thr, out_file=args.output)
+    show_result_pyplot(
+        model,
+        args.img,
+        result,
+        palette=args.palette,
+        score_thr=args.score_thr)
 
 
 if __name__ == '__main__':
-    main()
+    args = parse_args()
+    main(args)
diff --git a/docs/en/api.rst b/docs/en/api.rst
@@ -31,6 +31,11 @@ post_processing
 .. automodule:: mmrotate.core.post_processing
     :members:
 
+visualization
+^^^^^^^^^^^^^
+.. automodule:: mmrotate.core.visualization
+    :members:
+
 mmrotate.datasets
 --------------
 

diff --git a/docs/zh_cn/api.rst b/docs/zh_cn/api.rst
@@ -31,6 +31,11 @@ post_processing
 .. automodule:: mmrotate.core.post_processing
     :members:
 
+visualization
+^^^^^^^^^^^^^
+.. automodule:: mmrotate.core.visualization
+    :members:
+
 mmrotate.datasets
 --------------
 

diff --git a/mmrotate/core/visualization/__init__.py b/mmrotate/core/visualization/__init__.py
@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .image import imshow_det_rbboxes
+from .palette import get_palette
 
-__all__ = ['imshow_det_rbboxes']
+__all__ = ['imshow_det_rbboxes', 'get_palette']