Merge pull request #187 from xiaohu2015/master

support segmentation mask in COCO’s compressed RLE format for condinst
aim-uofa · Aug 21, 2020 · d459a2f · d459a2f
2 parents ba72974 + a666bf4
commit d459a2f
Showing 1 changed file with 28 additions and 18 deletions.
diff --git a/adet/modeling/condinst/condinst.py b/adet/modeling/condinst/condinst.py
@@ -139,24 +139,34 @@ def add_bitmasks(self, instances, im_h, im_w):
         for per_im_gt_inst in instances:
             if not per_im_gt_inst.has("gt_masks"):
                 continue
-            polygons = per_im_gt_inst.get("gt_masks").polygons
-            per_im_bitmasks = []
-            per_im_bitmasks_full = []
-            for per_polygons in polygons:
-                bitmask = polygons_to_bitmask(per_polygons, im_h, im_w)
-                bitmask = torch.from_numpy(bitmask).to(self.device).float()
-                start = int(self.mask_out_stride // 2)
-                bitmask_full = bitmask.clone()
-                bitmask = bitmask[start::self.mask_out_stride, start::self.mask_out_stride]
-
-                assert bitmask.size(0) * self.mask_out_stride == im_h
-                assert bitmask.size(1) * self.mask_out_stride == im_w
-
-                per_im_bitmasks.append(bitmask)
-                per_im_bitmasks_full.append(bitmask_full)
-
-            per_im_gt_inst.gt_bitmasks = torch.stack(per_im_bitmasks, dim=0)
-            per_im_gt_inst.gt_bitmasks_full = torch.stack(per_im_bitmasks_full, dim=0)
+            start = int(self.mask_out_stride // 2)
+            if isinstance(per_im_gt_inst.get("gt_masks"), PolygonMasks):
+                polygons = per_im_gt_inst.get("gt_masks").polygons
+                per_im_bitmasks = []
+                per_im_bitmasks_full = []
+                for per_polygons in polygons:
+                    bitmask = polygons_to_bitmask(per_polygons, im_h, im_w)
+                    bitmask = torch.from_numpy(bitmask).to(self.device).float()
+                    start = int(self.mask_out_stride // 2)
+                    bitmask_full = bitmask.clone()
+                    bitmask = bitmask[start::self.mask_out_stride, start::self.mask_out_stride]
+
+                    assert bitmask.size(0) * self.mask_out_stride == im_h
+                    assert bitmask.size(1) * self.mask_out_stride == im_w
+
+                    per_im_bitmasks.append(bitmask)
+                    per_im_bitmasks_full.append(bitmask_full)
+
+                per_im_gt_inst.gt_bitmasks = torch.stack(per_im_bitmasks, dim=0)
+                per_im_gt_inst.gt_bitmasks_full = torch.stack(per_im_bitmasks_full, dim=0)
+            else: # RLE format bitmask
+                bitmasks = per_im_gt_inst.get("gt_masks").tensor
+                h, w = bitmasks.size()[1:]
+                # pad to new size
+                bitmasks_full = F.pad(bitmasks, (0, im_w - w, 0, im_h - h), "constant", 0)
+                bitmasks = bitmasks_full[:, start::self.mask_out_stride, start::self.mask_out_stride]
+                per_im_gt_inst.gt_bitmasks = bitmasks
+                per_im_gt_inst.gt_bitmasks_full = bitmasks_full
 
     def postprocess(self, results, output_height, output_width, padded_im_h, padded_im_w, mask_threshold=0.5):
         """