Convert segmentation mask to binary mask (#55)

Author: michal-lightly

src/labelformat/model/semantic_segmentation.py

@@ -2,6 +2,9 @@
 
 from typing import List, Optional, Tuple
 
+from labelformat.model.binary_mask_segmentation import BinaryMaskSegmentation
+from labelformat.model.instance_segmentation import SingleInstanceSegmentation
+
 """Semantic segmentation core types and input interface.
 """
 
@@ -56,6 +59,28 @@ def from_array(cls, array: NDArray[np.int_]) -> "SemanticSegmentationMask":
             category_id_rle=category_id_rle, width=array.shape[1], height=array.shape[0]
         )
 
+    def to_binary_mask(self, category_id: int) -> BinaryMaskSegmentation:
+        """Get a binary mask for a given category ID."""
+        binary_rle = []
+
+        symbol = 0
+        run_length = 0
+        for cat_id, cur_run_length in self.category_id_rle:
+            cur_symbol = 1 if cat_id == category_id else 0
+            if symbol == cur_symbol:
+                run_length += cur_run_length
+            else:
+                binary_rle.append(run_length)
+                symbol = cur_symbol
+                run_length = cur_run_length
+
+        binary_rle.append(run_length)
+        return BinaryMaskSegmentation.from_rle(
+            rle_row_wise=binary_rle,
+            width=self.width,
+            height=self.height,
+        )
+
 
 class SemanticSegmentationInput(ABC):
 

tests/unit/model/test_semantic_segmentation.py

@@ -25,3 +25,35 @@ def test_from_array(self) -> None:
         assert mask.category_id_rle == expected_rle
         assert mask.width == 4
         assert mask.height == 3
+
+    def test_to_binary_mask(self) -> None:
+        mask = SemanticSegmentationMask.from_array(
+            array=np.array(
+                [
+                    [1, 1, 2, 2],
+                    [2, 1, 1, 1],
+                    [3, 3, 3, 3],
+                ],
+                dtype=np.int_,
+            )
+        )
+        binary_mask = mask.to_binary_mask(category_id=1)
+        assert binary_mask.get_binary_mask().tolist() == [
+            [1, 1, 0, 0],
+            [0, 1, 1, 1],
+            [0, 0, 0, 0],
+        ]
+
+        binary_mask = mask.to_binary_mask(category_id=2)
+        assert binary_mask.get_binary_mask().tolist() == [
+            [0, 0, 1, 1],
+            [1, 0, 0, 0],
+            [0, 0, 0, 0],
+        ]
+
+        binary_mask = mask.to_binary_mask(category_id=4)
+        assert binary_mask.get_binary_mask().tolist() == [
+            [0, 0, 0, 0],
+            [0, 0, 0, 0],
+            [0, 0, 0, 0],
+        ]