Improve default mask crop behavior (#32)

* Fix segment crops for objects fully inside the background ones

* Improve default behavior for CropCoveredSegments transform

* Improve type naming for polygon segments

Author: zhiltsov-max

datumaro/plugins/transforms.py

@@ -110,7 +110,9 @@ def crop_segments(
                     rle = mask_tools.mask_to_rle(s.image)
                 segments.append(rle)
 
-        segments = mask_tools.crop_covered_segments(segments, img_width, img_height)
+        segments = mask_tools.crop_covered_segments(
+            segments, img_width, img_height, ratio_tolerance=0
+        )
 
         new_anns = []
         for ann, new_segment in zip(segment_anns, segments):

datumaro/util/mask_tools.py

@@ -22,9 +22,10 @@ class CompressedRle(TypedDict):
 
 
 Rle = Union[CompressedRle, UncompressedRle]
-Polygon = List[List[int]]
+Polygon = List[int]
+PolygonGroup = List[Polygon]
 BboxCoords = NamedTuple("BboxCoords", [("x", int), ("y", int), ("w", int), ("h", int)])
-Segment = Union[Polygon, Rle]
+Segment = Union[PolygonGroup, Rle]
 
 BinaryMask = NewType("BinaryMask", np.ndarray)
 IndexMask = NewType("IndexMask", np.ndarray)
@@ -234,7 +235,7 @@ def is_uncompressed_rle(obj: Segment) -> bool:
     return isinstance(obj, dict) and isinstance(obj.get("counts"), bytes)
 
 
-def is_polygon(obj: Segment) -> bool:
+def is_polygon_group(obj: Segment) -> bool:
     return (
         isinstance(obj, list)
         and isinstance(obj[0], list)
@@ -259,7 +260,7 @@ def crop_covered_segments(
     ratio_tolerance: float = 0.001,
     area_threshold: int = 1,
     return_masks: bool = False,
-) -> List[Union[Optional[BinaryMask], Polygon]]:
+) -> List[Union[Optional[BinaryMask], List[Polygon]]]:
     """
     Find all segments occluded by others and crop them to the visible part only.
     Input segments are expected to be sorted from background to foreground.
@@ -314,13 +315,14 @@ def crop_covered_segments(
             area_top = sum(mask_utils.area(rle_top))
             area_ratio = area_top / area_bottom
 
-            # If a segment is already fully inside the top ones, stop accumulating the top
+            # If the top segment is (almost) fully inside the background one,
+            # we may need to skip it to avoid making a hole in the background object
             if abs(area_ratio - iou) < ratio_tolerance:
-                break
+                continue
 
             rles_top += rle_top
 
-        if not rles_top and is_polygon(wrapped_segments[i]) and not return_masks:
+        if not rles_top and is_polygon_group(wrapped_segments[i]) and not return_masks:
             output_segments.append(wrapped_segments[i])
             continue
 
@@ -334,7 +336,7 @@ def crop_covered_segments(
             bottom_mask -= top_mask
             bottom_mask[bottom_mask != 1] = 0
 
-        if not return_masks and is_polygon(wrapped_segments[i]):
+        if not return_masks and is_polygon_group(wrapped_segments[i]):
             output_segments.append(mask_to_polygons(bottom_mask, area_threshold=area_threshold))
         else:
             if np.sum(bottom_mask) < area_threshold:

tests/test_masks.py

@@ -115,6 +115,39 @@ def test_can_crop_covered_segments(self):
         for i, (e_mask, c_mask) in enumerate(zip(expected, computed)):
             self.assertTrue(np.array_equal(e_mask, c_mask), "#%s: %s\n%s\n" % (i, e_mask, c_mask))
 
+    @mark_requirement(Requirements.DATUM_GENERAL_REQ)
+    def test_can_crop_covered_segments_and_avoid_holes_from_objects_inside_background_object(self):
+        image_size = [7, 7]
+        initial = [
+            [1, 1, 6, 1, 6, 6, 1, 6],
+            mask_tools.mask_to_rle(
+                np.array(
+                    [
+                        [0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 1, 1, 1, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0],
+                    ]
+                )
+            ),
+        ]
+        expected = [
+            # no changes expected
+            mask_tools.rles_to_mask([initial[0]], *image_size),
+            mask_tools.rles_to_mask([initial[1]], *image_size),
+        ]
+
+        computed = mask_tools.crop_covered_segments(
+            initial, *image_size, ratio_tolerance=0.1, return_masks=True
+        )
+
+        self.assertEqual(len(initial), len(computed))
+        for i, (e_mask, c_mask) in enumerate(zip(expected, computed)):
+            self.assertTrue(np.array_equal(e_mask, c_mask), "#%s: %s\n%s\n" % (i, e_mask, c_mask))
+
     def _test_mask_to_rle(self, source_mask):
         rle_uncompressed = mask_tools.mask_to_rle(source_mask)