rectify the position of masker

(cherry picked from commit ee91261)

demo/predictor.py

@@ -100,6 +100,7 @@ class COCODemo(object):
     def __init__(
             self,
             cfg,
+            masker,
             confidence_threshold=0.7,
             show_mask_heatmaps=False,
             masks_per_dim=2,
@@ -118,8 +119,10 @@ def __init__(
 
         self.transforms = self.build_transform()
 
-        mask_threshold = -1 if show_mask_heatmaps else 0.5
-        self.masker = Masker(threshold=mask_threshold, padding=1)
+        if show_mask_heatmaps:
+            self.masker = Masker(threshold=-1, padding=1)
+        else:
+            self.masker = masker
 
         # used to make colors for each class
         self.palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
@@ -217,7 +220,7 @@ def compute_prediction(self, original_image):
             # in the image, as defined by the bounding boxes
             masks = prediction.get_field("mask")
             # always single image is passed at a time
-            masks = self.masker([masks], [prediction])[0]
+            masks = self.masker.forward_single_image(masks, prediction)
             prediction.add_field("mask", masks)
         return prediction
 

maskrcnn_benchmark/data/datasets/evaluation/coco/__init__.py

@@ -9,6 +9,7 @@ def coco_evaluation(
     iou_types,
     expected_results,
     expected_results_sigma_tol,
+    masker,
 ):
     return do_coco_evaluation(
         dataset=dataset,
@@ -18,4 +19,5 @@ def coco_evaluation(
         iou_types=iou_types,
         expected_results=expected_results,
         expected_results_sigma_tol=expected_results_sigma_tol,
+        masker=masker,
     )

maskrcnn_benchmark/data/datasets/evaluation/coco/coco_eval.py

@@ -19,6 +19,7 @@ def do_coco_evaluation(
         iou_types,
         expected_results,
         expected_results_sigma_tol,
+        masker,
 ):
     logger = logging.getLogger("maskrcnn_benchmark.inference")
 
@@ -45,7 +46,7 @@ def do_coco_evaluation(
         coco_results["bbox"] = prepare_for_coco_detection(predictions, dataset)
     if "segm" in iou_types:
         logger.info("Preparing segm results")
-        coco_results["segm"] = prepare_for_coco_segmentation(predictions, dataset)
+        coco_results["segm"] = prepare_for_coco_segmentation(predictions, dataset, masker)
     if 'keypoints' in iou_types:
         logger.info('Preparing keypoints results')
         coco_results['keypoints'] = prepare_for_coco_keypoint(predictions, dataset)
@@ -102,11 +103,11 @@ def prepare_for_coco_detection(predictions, dataset):
     return coco_results
 
 
-def prepare_for_coco_segmentation(predictions, dataset):
+def prepare_for_coco_segmentation(predictions, dataset, masker):
     import pycocotools.mask as mask_util
     import numpy as np
 
-    masker = Masker(threshold=0.5, padding=1)
+    assert isinstance(masker, Masker)
     # assert isinstance(dataset, COCODataset)
     coco_results = []
     for image_id, prediction in tqdm(enumerate(predictions)):
@@ -120,10 +121,7 @@ def prepare_for_coco_segmentation(predictions, dataset):
         prediction = prediction.resize((image_width, image_height))
         masks = prediction.get_field("mask")
         # t = time.time()
-        # Masker is necessary only if masks haven't been already resized.
-        if list(masks.shape[-2:]) != [image_height, image_width]:
-            masks = masker(masks.expand(1, -1, -1, -1, -1), prediction)
-            masks = masks[0]
+        masks = masker.forward_single_image(masks, prediction)
         # logger.info('Time mask: {}'.format(time.time() - t))
         # prediction = prediction.convert('xywh')
 

maskrcnn_benchmark/engine/inference.py

@@ -6,6 +6,7 @@
 from tqdm import tqdm
 
 from maskrcnn_benchmark.data.datasets.evaluation import evaluate
+from maskrcnn_benchmark.modeling.roi_heads.mask_head.inference import Masker
 from ..utils.comm import all_gather
 from ..utils.comm import is_main_process, get_world_size
 from ..utils.comm import synchronize
@@ -59,6 +60,7 @@ def inference(
         model,
         data_loader,
         dataset_name,
+        masker,
         iou_types=("bbox",),
         box_only=False,
         device="cuda",
@@ -101,11 +103,14 @@ def inference(
     if output_folder:
         torch.save(predictions, os.path.join(output_folder, "predictions.pth"))
 
+    assert isinstance(masker, Masker)
+
     extra_args = dict(
         box_only=box_only,
         iou_types=iou_types,
         expected_results=expected_results,
         expected_results_sigma_tol=expected_results_sigma_tol,
+        masker=masker
     )
 
     return evaluate(dataset=dataset,

maskrcnn_benchmark/modeling/roi_heads/mask_head/inference.py

@@ -200,10 +200,10 @@ def __call__(self, masks, boxes):
 
 
 def make_roi_mask_post_processor(cfg):
+    # In the statement "mask_prob = self.masker(mask_prob, boxes)" from MaskPostProcessor.forward
+    # the image size got from boxes is the resized image shape (INPUT.MAX_SIZE_TEST), not the origin image shape
+    # even doing postprocess in this stage, we still need to "masker" again in the coco_eval method
+    # so, we decide to delete this masker, and move it to the evaluation method
     if cfg.MODEL.ROI_MASK_HEAD.POSTPROCESS_MASKS:
-        mask_threshold = cfg.MODEL.ROI_MASK_HEAD.POSTPROCESS_MASKS_THRESHOLD
-        masker = Masker(threshold=mask_threshold, padding=1)
-    else:
-        masker = None
-    mask_post_processor = MaskPostProcessor(masker)
-    return mask_post_processor
+        raise ValueError("Not recommend doing postprocess in this stage, see the above explanation for details.")
+    return MaskPostProcessor(masker=None)

tools/test_net.py

@@ -11,6 +11,7 @@
 from maskrcnn_benchmark.data import make_data_loader
 from maskrcnn_benchmark.engine.inference import inference
 from maskrcnn_benchmark.modeling.detector import build_detection_model
+from maskrcnn_benchmark.modeling.roi_heads.mask_head.inference import Masker
 from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
 from maskrcnn_benchmark.utils.collect_env import collect_env_info
 from maskrcnn_benchmark.utils.comm import synchronize, get_rank
@@ -88,11 +89,13 @@ def main():
             mkdir(output_folder)
             output_folders[idx] = output_folder
     data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)
+    masker = Masker(threshold=0.5, padding=1)
     for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val):
         inference(
             model,
             data_loader_val,
             dataset_name=dataset_name,
+            masker=masker,
             iou_types=iou_types,
             box_only=False if cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY,
             device=cfg.MODEL.DEVICE,