From 97a712ee445292806ba63c8306ce041fdf6b60bc Mon Sep 17 00:00:00 2001
From: Cedric Luo <luochunhua1996@outlook.com>
Date: Thu, 24 Feb 2022 23:02:50 +0800
Subject: [PATCH] Maskformer Visualization (#7247)

* maskformer visualization

* compatible with models that do not contain arg of pretrained

* compatible with models that do not contain arg of pretrained
---
 mmdet/apis/inference.py              |  5 +-
 mmdet/models/detectors/maskformer.py | 83 ++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/mmdet/apis/inference.py b/mmdet/apis/inference.py
index 6d58d36ff1a..393037cf99d 100644
--- a/mmdet/apis/inference.py
+++ b/mmdet/apis/inference.py
@@ -34,7 +34,10 @@ def init_detector(config, checkpoint=None, device='cuda:0', cfg_options=None):
                         f'but got {type(config)}')
     if cfg_options is not None:
         config.merge_from_dict(cfg_options)
-    config.model.pretrained = None
+    if 'pretrained' in config.model:
+        config.model.pretrained = None
+    elif 'init_cfg' in config.model.backbone:
+        config.model.backbone.init_cfg = None
     config.model.train_cfg = None
     model = build_detector(config.model, test_cfg=config.get('test_cfg'))
     if checkpoint is not None:
diff --git a/mmdet/models/detectors/maskformer.py b/mmdet/models/detectors/maskformer.py
index 73676bcdf50..f7257d2547d 100644
--- a/mmdet/models/detectors/maskformer.py
+++ b/mmdet/models/detectors/maskformer.py
@@ -1,4 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import mmcv
+import numpy as np
+
+from mmdet.core import INSTANCE_OFFSET
+from mmdet.core.visualization import imshow_det_bboxes
 from ..builder import DETECTORS, build_backbone, build_head, build_neck
 from .single_stage import SingleStageDetector
 
@@ -23,6 +28,11 @@ def __init__(self,
         panoptic_head.update(train_cfg=train_cfg)
         panoptic_head.update(test_cfg=test_cfg)
         self.panoptic_head = build_head(panoptic_head)
+
+        self.num_things_classes = self.panoptic_head.num_things_classes
+        self.num_stuff_classes = self.panoptic_head.num_stuff_classes
+        self.num_classes = self.panoptic_head.num_classes
+
         self.train_cfg = train_cfg
         self.test_cfg = test_cfg
 
@@ -104,3 +114,76 @@ def aug_test(self, imgs, img_metas, **kwargs):
 
     def onnx_export(self, img, img_metas):
         raise NotImplementedError
+
+    def show_result(self,
+                    img,
+                    result,
+                    score_thr=0.3,
+                    bbox_color=(72, 101, 241),
+                    text_color=(72, 101, 241),
+                    mask_color=None,
+                    thickness=2,
+                    font_size=13,
+                    win_name='',
+                    show=False,
+                    wait_time=0,
+                    out_file=None):
+        """Draw `result` over `img`.
+
+        Args:
+            img (str or Tensor): The image to be displayed.
+            result (dict): The results.
+
+            score_thr (float, optional): Minimum score of bboxes to be shown.
+                Default: 0.3.
+            bbox_color (str or tuple(int) or :obj:`Color`):Color of bbox lines.
+               The tuple of color should be in BGR order. Default: 'green'.
+            text_color (str or tuple(int) or :obj:`Color`):Color of texts.
+               The tuple of color should be in BGR order. Default: 'green'.
+            mask_color (None or str or tuple(int) or :obj:`Color`):
+               Color of masks. The tuple of color should be in BGR order.
+               Default: None.
+            thickness (int): Thickness of lines. Default: 2.
+            font_size (int): Font size of texts. Default: 13.
+            win_name (str): The window name. Default: ''.
+            wait_time (float): Value of waitKey param.
+                Default: 0.
+            show (bool): Whether to show the image.
+                Default: False.
+            out_file (str or None): The filename to write the image.
+                Default: None.
+
+        Returns:
+            img (Tensor): Only if not `show` or `out_file`.
+        """
+        img = mmcv.imread(img)
+        img = img.copy()
+        pan_results = result['pan_results']
+        # keep objects ahead
+        ids = np.unique(pan_results)[::-1]
+        legal_indices = ids != self.num_classes  # for VOID label
+        ids = ids[legal_indices]
+        labels = np.array([id % INSTANCE_OFFSET for id in ids], dtype=np.int64)
+        segms = (pan_results[None] == ids[:, None, None])
+
+        # if out_file specified, do not show image in window
+        if out_file is not None:
+            show = False
+        # draw bounding boxes
+        img = imshow_det_bboxes(
+            img,
+            segms=segms,
+            labels=labels,
+            class_names=self.CLASSES,
+            bbox_color=bbox_color,
+            text_color=text_color,
+            mask_color=mask_color,
+            thickness=thickness,
+            font_size=font_size,
+            win_name=win_name,
+            show=show,
+            wait_time=wait_time,
+            out_file=out_file)
+
+        if not (show or out_file):
+            return img