handle overlapped segmentations more gracefully (#849)

Author: czaloom

src/valor_lite/semantic_segmentation/annotation.py

@@ -1,3 +1,4 @@
+import warnings
 from dataclasses import dataclass, field
 
 import numpy as np
@@ -79,30 +80,37 @@ def __post_init__(self):
             )
         self.size = self.shape[0] * self.shape[1]
 
-        mask_accumulation = None
-        for groundtruth in self.groundtruths:
-            if self.shape != groundtruth.mask.shape:
-                raise ValueError(
-                    f"ground truth masks for datum '{self.uid}' should have shape '{self.shape}'. Received mask with shape '{groundtruth.mask.shape}'"
-                )
-
-            if mask_accumulation is None:
-                mask_accumulation = groundtruth.mask.copy()
-            elif np.logical_and(mask_accumulation, groundtruth.mask).any():
-                raise ValueError("ground truth masks cannot overlap")
-            else:
-                mask_accumulation = mask_accumulation | groundtruth.mask
+        self._validate_bitmasks(self.groundtruths, "ground truth")
+        self._validate_bitmasks(self.predictions, "prediction")
 
+    def _validate_bitmasks(self, bitmasks: list[Bitmask], key: str):
         mask_accumulation = None
-        for prediction in self.predictions:
-            if self.shape != prediction.mask.shape:
+        mask_overlap_accumulation = None
+        for idx, bitmask in enumerate(bitmasks):
+            if not isinstance(bitmask, Bitmask):
+                raise ValueError(f"expected 'Bitmask', got '{bitmask}'")
+            if self.shape != bitmask.mask.shape:
                 raise ValueError(
-                    f"prediction masks for datum '{self.uid}' should have shape '{self.shape}'. Received mask with shape '{prediction.mask.shape}'"
+                    f"{key} masks for datum '{self.uid}' should have shape '{self.shape}'. Received mask with shape '{bitmask.mask.shape}'"
                 )
 
             if mask_accumulation is None:
-                mask_accumulation = prediction.mask.copy()
-            elif np.logical_and(mask_accumulation, prediction.mask).any():
-                raise ValueError("prediction masks cannot overlap")
+                mask_accumulation = bitmask.mask.copy()
+                mask_overlap_accumulation = np.zeros_like(mask_accumulation)
+            elif np.logical_and(mask_accumulation, bitmask.mask).any():
+                mask_overlap = np.logical_and(mask_accumulation, bitmask.mask)
+                bitmasks[idx].mask[mask_overlap] = False
+                mask_overlap_accumulation = (
+                    mask_overlap_accumulation | mask_overlap
+                )
             else:
-                mask_accumulation = mask_accumulation | prediction.mask
+                mask_accumulation = mask_accumulation | bitmask.mask
+        if (
+            mask_overlap_accumulation is not None
+            and mask_overlap_accumulation.any()
+        ):
+            count = mask_overlap_accumulation.sum()
+            total = mask_overlap_accumulation.size
+            warnings.warn(
+                f"{key} masks for datum '{self.uid}' had {count} / {total} pixels overlapped."
+            )

tests/semantic_segmentation/test_annotation.py

@@ -105,54 +105,6 @@ def test_segmentation():
         )
     assert "Received mask with shape '(1, 2)'" in str(e)
 
-    # test ground truths cannot overlap
-    with pytest.raises(ValueError) as e:
-        Segmentation(
-            uid="uid",
-            groundtruths=[
-                Bitmask(
-                    mask=np.array([[True, True, True]]),
-                    label="label1",
-                ),
-                Bitmask(
-                    mask=np.array([[False, False, True]]),
-                    label="label2",
-                ),
-            ],
-            predictions=[
-                Bitmask(
-                    mask=np.array([[True, False, False]]),
-                    label="label",
-                )
-            ],
-            shape=(1, 3),
-        )
-    assert "ground truth masks cannot overlap" in str(e)
-
-    # test predictions cannot overlap
-    with pytest.raises(ValueError) as e:
-        Segmentation(
-            uid="uid",
-            groundtruths=[
-                Bitmask(
-                    mask=np.array([[True, True, True]]),
-                    label="label1",
-                ),
-            ],
-            predictions=[
-                Bitmask(
-                    mask=np.array([[True, False, True]]),
-                    label="label",
-                ),
-                Bitmask(
-                    mask=np.array([[False, False, True]]),
-                    label="label2",
-                ),
-            ],
-            shape=(1, 3),
-        )
-    assert "prediction masks cannot overlap" in str(e)
-
     # allow missing ground truths
     Segmentation(
         uid="uid",
@@ -179,6 +131,22 @@ def test_segmentation():
         shape=(1, 2),
     )
 
+    # wrong annotation type
+    with pytest.raises(ValueError):
+        Segmentation(
+            uid="uid",
+            groundtruths=[{"a": 1}],  # type: ignore - testing
+            predictions=[],
+            shape=(1, 2),
+        )
+    with pytest.raises(ValueError):
+        Segmentation(
+            uid="uid",
+            groundtruths=[],
+            predictions=[{"a": 1}],  # type: ignore - testing
+            shape=(1, 2),
+        )
+
 
 def test_segmentation_shape():
     Segmentation(uid="uid", groundtruths=[], predictions=[], shape=(1, 1))
@@ -197,3 +165,58 @@ def test_segmentation_shape():
         Segmentation(
             uid="uid", groundtruths=[], predictions=[], shape=(-100, 100)
         )
+
+
+def _create_overlapped_masks() -> tuple[Bitmask, Bitmask]:
+    mask0 = np.zeros((100, 100), dtype=np.bool_)
+    mask0[:50, :] = True
+    mask1 = np.ones((100, 100), dtype=np.bool_)
+    bitmask0 = Bitmask(mask=mask0, label="dog")
+    bitmask1 = Bitmask(mask=mask1, label="cat")
+    return bitmask0, bitmask1
+
+
+def test_segmentations_overlap():
+    bitmask0, bitmask1 = _create_overlapped_masks()
+    Segmentation(
+        uid="uid123",
+        groundtruths=[bitmask0],
+        predictions=[bitmask1],
+        shape=(100, 100),
+    )
+    assert bitmask0.mask.sum() == 5000
+    assert bitmask1.mask.sum() == 10000
+
+    bitmask0, bitmask1 = _create_overlapped_masks()
+    with pytest.warns(UserWarning) as e:
+        Segmentation(
+            uid="uid123",
+            groundtruths=[bitmask0, bitmask1],
+            predictions=[],
+            shape=(100, 100),
+        )
+    assert (
+        str(e._list[0].message)
+        == "ground truth masks for datum 'uid123' had 5000 / 10000 pixels overlapped."
+    )
+    assert bitmask0.mask.sum() == 5000
+    assert (
+        bitmask1.mask.sum() == 5000
+    )  # overlapped pixels omitted from second mask
+
+    bitmask0, bitmask1 = _create_overlapped_masks()
+    with pytest.warns(UserWarning) as e:
+        Segmentation(
+            uid="uid123",
+            groundtruths=[],
+            predictions=[bitmask0, bitmask1],
+            shape=(100, 100),
+        )
+    assert (
+        str(e._list[0].message)
+        == "prediction masks for datum 'uid123' had 5000 / 10000 pixels overlapped."
+    )
+    assert bitmask0.mask.sum() == 5000
+    assert (
+        bitmask1.mask.sum() == 5000
+    )  # overlapped pixels omitted from second mask