fix show bug (#2824)

mmdet/apis/test.py

@@ -4,34 +4,11 @@ import shutil
 import tempfile
 
 import mmcv
-import numpy as np
-import pycocotools.mask as mask_util
 import torch
 import torch.distributed as dist
 from mmcv.runner import get_dist_info
 
-from mmdet.core import tensor2imgs
-
-
-# TODO: move this function to more proper place
-def encode_mask_results(mask_results):
-    if isinstance(mask_results, tuple):  # mask scoring
-        cls_segms, cls_mask_scores = mask_results
-    else:
-        cls_segms = mask_results
-    num_classes = len(cls_segms)
-    encoded_mask_results = [[] for _ in range(num_classes)]
-    for i in range(len(cls_segms)):
-        for cls_segm in cls_segms[i]:
-            encoded_mask_results[i].append(
-                mask_util.encode(
-                    np.array(
-                        cls_segm[:, :, np.newaxis], order='F',
-                        dtype='uint8'))[0])  # encoded with RLE
-    if isinstance(mask_results, tuple):
-        return encoded_mask_results, cls_mask_scores
-    else:
-        return encoded_mask_results
+from mmdet.core import encode_mask_results, tensor2imgs
 
 
 def single_gpu_test(model,
@@ -46,12 +23,6 @@ def single_gpu_test(model,
     for i, data in enumerate(data_loader):
         with torch.no_grad():
             result = model(return_loss=False, rescale=True, **data)
-            # encode mask results
-            if isinstance(result, tuple):
-                bbox_results, mask_results = result
-                encoded_mask_results = encode_mask_results(mask_results)
-                result = bbox_results, encoded_mask_results
-        results.append(result)
 
         if show or out_dir:
             img_tensor = data['img'][0]
@@ -78,6 +49,13 @@ def single_gpu_test(model,
                     out_file=out_file,
                     score_thr=show_score_thr)
 
+        # encode mask results
+        if isinstance(result, tuple):
+            bbox_results, mask_results = result
+            encoded_mask_results = encode_mask_results(mask_results)
+            result = bbox_results, encoded_mask_results
+        results.append(result)
+
         batch_size = data['img'][0].size(0)
         for _ in range(batch_size):
             prog_bar.update()

mmdet/core/mask/__init__.py

@@ -1,7 +1,8 @@
 from .mask_target import mask_target
 from .structures import BitmapMasks, PolygonMasks
-from .utils import split_combined_polys
+from .utils import encode_mask_results, split_combined_polys
 
 __all__ = [
-    'split_combined_polys', 'mask_target', 'BitmapMasks', 'PolygonMasks'
+    'split_combined_polys', 'mask_target', 'BitmapMasks', 'PolygonMasks',
+    'encode_mask_results'
 ]

mmdet/core/mask/utils.py

@@ -1,4 +1,6 @@
 import mmcv
+import numpy as np
+import pycocotools.mask as mask_util
 
 
 def split_combined_polys(polys, poly_lens, polys_per_mask):
@@ -28,3 +30,34 @@ def split_combined_polys(polys, poly_lens, polys_per_mask):
         mask_polys = mmcv.slice_list(split_polys, polys_per_mask_single)
         mask_polys_list.append(mask_polys)
     return mask_polys_list
+
+
+# TODO: move this function to more proper place
+def encode_mask_results(mask_results):
+    """Encode bitmap mask to RLE code.
+
+    Args:
+        mask_results (list | tuple[list]): bitmap mask results.
+            In mask scoring rcnn, mask_results is a tuple of (segm_results,
+            segm_cls_score).
+
+    Returns:
+        list | tuple: RLE encoded mask.
+    """
+    if isinstance(mask_results, tuple):  # mask scoring
+        cls_segms, cls_mask_scores = mask_results
+    else:
+        cls_segms = mask_results
+    num_classes = len(cls_segms)
+    encoded_mask_results = [[] for _ in range(num_classes)]
+    for i in range(len(cls_segms)):
+        for cls_segm in cls_segms[i]:
+            encoded_mask_results[i].append(
+                mask_util.encode(
+                    np.array(
+                        cls_segm[:, :, np.newaxis], order='F',
+                        dtype='uint8'))[0])  # encoded with RLE
+    if isinstance(mask_results, tuple):
+        return encoded_mask_results, cls_mask_scores
+    else:
+        return encoded_mask_results

mmdet/models/detectors/base.py

@@ -3,7 +3,6 @@ from abc import ABCMeta, abstractmethod
 
 import mmcv
 import numpy as np
-import pycocotools.mask as maskUtils
 import torch.nn as nn
 from mmcv.utils import print_log
 
@@ -210,7 +209,7 @@ class BaseDetector(nn.Module, metaclass=ABCMeta):
             for i in inds:
                 i = int(i)
                 color_mask = color_masks[labels[i]]
-                mask = maskUtils.decode(segms[i]).astype(np.bool)
+                mask = segms[i]
                 img[mask] = img[mask] * 0.5 + color_mask * 0.5
         # if out_file specified, do not show image in window
         if out_file is not None:

mmdet/models/detectors/cascade_rcnn.py

@@ -31,4 +31,4 @@ class CascadeRCNN(TwoStageDetector):
         else:
             if isinstance(result, dict):
                 result = result['ensemble']
-        super(CascadeRCNN, self).show_result(data, result, **kwargs)
+        return super(CascadeRCNN, self).show_result(data, result, **kwargs)

tools/test_robustness.py

@@ -16,7 +16,7 @@ from robustness_eval import get_results
 
 from mmdet import datasets
 from mmdet.apis import set_random_seed
-from mmdet.core import eval_map, wrap_fp16_model
+from mmdet.core import encode_mask_results, eval_map, wrap_fp16_model
 from mmdet.datasets import build_dataloader, build_dataset
 from mmdet.models import build_detector
 
@@ -98,11 +98,17 @@ def single_gpu_test(model, data_loader, show=False):
     for i, data in enumerate(data_loader):
         with torch.no_grad():
             result = model(return_loss=False, rescale=not show, **data)
-        results.append(result)
 
         if show:
             model.module.show_result(data, result, dataset.img_norm_cfg)
 
+        # encode mask results
+        if isinstance(result, tuple):
+            bbox_results, mask_results = result
+            encoded_mask_results = encode_mask_results(mask_results)
+            result = bbox_results, encoded_mask_results
+        results.append(result)
+
         batch_size = data['img'][0].size(0)
         for _ in range(batch_size):
             prog_bar.update()
@@ -119,6 +125,13 @@ def multi_gpu_test(model, data_loader, tmpdir=None):
     for i, data in enumerate(data_loader):
         with torch.no_grad():
             result = model(return_loss=False, rescale=True, **data)
+            # encode mask results
+            if isinstance(result, tuple):
+                bbox_results, mask_results = result
+                encoded_mask_results = encode_mask_results(mask_results)
+                result = bbox_results, encoded_mask_results
+        results.append(result)
+
         results.append(result)
 
         if rank == 0: