From f0ec12dd71f443eaebd885a81bde0368d027498e Mon Sep 17 00:00:00 2001 From: Cao Yuhang Date: Wed, 27 May 2020 11:39:56 +0800 Subject: [PATCH] fix show bug (#2824) --- mmdet/apis/test.py | 38 ++++++-------------------- mmdet/core/mask/__init__.py | 5 ++-- mmdet/core/mask/utils.py | 33 ++++++++++++++++++++++ mmdet/models/detectors/base.py | 3 +- mmdet/models/detectors/cascade_rcnn.py | 2 +- tools/test_robustness.py | 17 ++++++++++-- 6 files changed, 61 insertions(+), 37 deletions(-) diff --git a/mmdet/apis/test.py b/mmdet/apis/test.py index 4d28053b2..9e7e1c234 100644 --- a/mmdet/apis/test.py +++ b/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() diff --git a/mmdet/core/mask/__init__.py b/mmdet/core/mask/__init__.py index ed7292111..4d8e9da19 100644 --- a/mmdet/core/mask/__init__.py +++ b/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' ] diff --git a/mmdet/core/mask/utils.py b/mmdet/core/mask/utils.py index a68312b17..34e60b25a 100644 --- a/mmdet/core/mask/utils.py +++ b/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 diff --git a/mmdet/models/detectors/base.py b/mmdet/models/detectors/base.py index 05236371c..e0cfb603b 100644 --- a/mmdet/models/detectors/base.py +++ b/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: diff --git a/mmdet/models/detectors/cascade_rcnn.py b/mmdet/models/detectors/cascade_rcnn.py index 367e7fe5a..d3475e598 100644 --- a/mmdet/models/detectors/cascade_rcnn.py +++ b/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) diff --git a/tools/test_robustness.py b/tools/test_robustness.py index ecc366f54..ca308522d 100644 --- a/tools/test_robustness.py +++ b/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: