Support Binary Mask with transparent SementationMask interface (#473)

* support RLE and binary mask * do not convert to numpy * be consistent with Detectron * delete wrong comment * [WIP] add tests for segmentation_mask * update tests * minor change * Refactored segmentation_mask.py * Add unit test for segmentation_mask.py * Add RLE support for BinaryMaskList * PEP8 black formatting * Minor patch * Use internal that handles 0 channels * Fix polygon slicing
facebookresearch · Apr 9, 2019 · b4d5465 · b4d5465 · zimenglan-sysu-512 · Apr 18, 2019
1 parent f917a55
commit b4d5465
Show file tree

Hide file tree

Showing 4 changed files with 478 additions and 86 deletions.
diff --git a/maskrcnn_benchmark/data/datasets/coco.py b/maskrcnn_benchmark/data/datasets/coco.py
@@ -80,7 +80,7 @@ def __getitem__(self, idx):
         target.add_field("labels", classes)
 
         masks = [obj["segmentation"] for obj in anno]
-        masks = SegmentationMask(masks, img.size)
+        masks = SegmentationMask(masks, img.size, mode='poly')
         target.add_field("masks", masks)
 
         if anno and "keypoints" in anno[0]:

diff --git a/maskrcnn_benchmark/modeling/roi_heads/mask_head/loss.py b/maskrcnn_benchmark/modeling/roi_heads/mask_head/loss.py
@@ -27,17 +27,15 @@ def project_masks_on_boxes(segmentation_masks, proposals, discretization_size):
     assert segmentation_masks.size == proposals.size, "{}, {}".format(
         segmentation_masks, proposals
     )
-    # TODO put the proposals on the CPU, as the representation for the
-    # masks is not efficient GPU-wise (possibly several small tensors for
-    # representing a single instance mask)
+
+    # FIXME: CPU computation bottleneck, this should be parallelized
     proposals = proposals.bbox.to(torch.device("cpu"))
     for segmentation_mask, proposal in zip(segmentation_masks, proposals):
         # crop the masks, resize them to the desired resolution and
-        # then convert them to the tensor representation,
-        # instead of the list representation that was used
+        # then convert them to the tensor representation.
         cropped_mask = segmentation_mask.crop(proposal)
         scaled_mask = cropped_mask.resize((M, M))
-        mask = scaled_mask.convert(mode="mask")
+        mask = scaled_mask.get_mask_tensor()
         masks.append(mask)
     if len(masks) == 0:
         return torch.empty(0, dtype=torch.float32, device=device)