Create binary segmentation mask from RLE (#54)

Author: michal-lightly

src/labelformat/model/binary_mask_segmentation.py

@@ -42,6 +42,28 @@ def from_binary_mask(
             bounding_box=bounding_box,
         )
 
+    @classmethod
+    def from_rle(
+        cls,
+        rle_row_wise: list[int],
+        width: int,
+        height: int,
+        bounding_box: BoundingBox | None = None,
+    ) -> "BinaryMaskSegmentation":
+        """
+        Create a BinaryMaskSegmentation instance from row-wise RLE format.
+        """
+        if bounding_box is None:
+            bounding_box = _compute_bbox_from_rle(
+                rle_row_wise=rle_row_wise, width=width, height=height
+            )
+        return cls(
+            _rle_row_wise=rle_row_wise,
+            width=width,
+            height=height,
+            bounding_box=bounding_box,
+        )
+
     def get_binary_mask(self) -> NDArray[np.int_]:
         """
         Get the binary mask (2D numpy array) from the RLE format.
@@ -112,3 +134,49 @@ def decode_column_wise_rle(
             decoded.extend([run_val] * count)
             run_val = 1 - run_val
         return np.array(decoded, dtype=np.int_).reshape((height, width), order="F")
+
+
+def _compute_bbox_from_rle(
+    rle_row_wise: list[int], width: int, height: int
+) -> BoundingBox:
+    """Compute bounding box from row-wise RLE.
+
+    Scans through the RLE and tracks the min/max x/y coordinates of the '1' pixels.
+    The time complexity is O(len(rle_row_wise)).
+    """
+    xmin = width
+    ymin = height
+    xmax = 0
+    ymax = 0
+
+    x = 0
+    y = 0
+    next_symbol = 0
+    for run_length in rle_row_wise:
+        if next_symbol == 1:
+            # Compute coordinates for the end of the run
+            run_end_x = x + run_length - 1
+            run_end_y = y
+            if run_end_x >= width:
+                run_end_y += run_end_x // width
+                run_end_x = run_end_x % width
+
+            # Update bounding box
+            ymin = min(ymin, y)
+            ymax = max(ymax, run_end_y)
+            if run_end_y > y:
+                xmin = 0
+                xmax = width - 1
+            else:
+                xmin = min(xmin, x)
+                xmax = max(xmax, run_end_x)
+
+        # Compute coordinates for the start of the next run
+        x += run_length
+        if x >= width:
+            y += x // width
+            x = x % width
+
+        next_symbol = 1 - next_symbol
+
+    return BoundingBox(xmin=xmin, ymin=ymin, xmax=xmax, ymax=ymax)

tests/unit/model/test_binary_mask_segmentation.py

@@ -1,6 +1,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from labelformat.model import binary_mask_segmentation
 from labelformat.model.binary_mask_segmentation import (
     BinaryMaskSegmentation,
     RLEDecoderEncoder,
@@ -22,6 +23,48 @@ def test_from_binary_mask(self) -> None:
         assert binary_mask_segmentation.bounding_box == bounding_box
         assert np.array_equal(binary_mask_segmentation.get_binary_mask(), binary_mask)
 
+    def test_from_rle(self) -> None:
+        binary_mask_segmentation = BinaryMaskSegmentation.from_rle(
+            rle_row_wise=[1, 1, 4, 2, 1, 3, 2, 1, 5],
+            width=5,
+            height=4,
+            bounding_box=None,
+        )
+        assert binary_mask_segmentation.width == 5
+        assert binary_mask_segmentation.height == 4
+        assert binary_mask_segmentation.bounding_box == BoundingBox(0, 0, 4, 2)
+        expected: NDArray[np.int_] = np.array(
+            [
+                [0, 1, 0, 0, 0],
+                [0, 1, 1, 0, 1],
+                [1, 1, 0, 0, 1],
+                [0, 0, 0, 0, 0],
+            ],
+            dtype=np.int_,
+        )
+        assert np.array_equal(binary_mask_segmentation.get_binary_mask(), expected)
+
+        # Test with provided bounding box
+        # The box is larger than the actual mask, but should be preserved
+        binary_mask_segmentation = BinaryMaskSegmentation.from_rle(
+            rle_row_wise=[6, 3],
+            width=3,
+            height=3,
+            bounding_box=BoundingBox(0, 0, 2, 2),
+        )
+        assert binary_mask_segmentation.width == 3
+        assert binary_mask_segmentation.height == 3
+        assert binary_mask_segmentation.bounding_box == BoundingBox(0, 0, 2, 2)
+        expected = np.array(
+            [
+                [0, 0, 0],
+                [0, 0, 0],
+                [1, 1, 1],
+            ],
+            dtype=np.int_,
+        )
+        assert np.array_equal(binary_mask_segmentation.get_binary_mask(), expected)
+
 
 class TestRLEDecoderEncoder:
     def test_encode_row_wise_rle(self) -> None:
@@ -79,3 +122,41 @@ def test_inverse__column_wise(self) -> None:
             rle, mask.shape[0], mask.shape[1]
         )
         assert np.array_equal(mask, mask_inverse_column_wise)
+
+
+def test_compute_bbox_from_rle() -> None:
+    # 0011
+    # 1111
+    # 1100
+    bbox = binary_mask_segmentation._compute_bbox_from_rle(
+        rle_row_wise=[2, 8, 2],
+        width=4,
+        height=3,
+    )
+    assert bbox == BoundingBox(xmin=0, ymin=0, xmax=3, ymax=2)
+
+    # 0011
+    # 0000
+    bbox = binary_mask_segmentation._compute_bbox_from_rle(
+        rle_row_wise=[2, 2, 4],
+        width=4,
+        height=2,
+    )
+    assert bbox == BoundingBox(xmin=2, ymin=0, xmax=3, ymax=0)
+
+    # 0011
+    # 1000
+    bbox = binary_mask_segmentation._compute_bbox_from_rle(
+        rle_row_wise=[2, 3, 3],
+        width=4,
+        height=2,
+    )
+    assert bbox == BoundingBox(xmin=0, ymin=0, xmax=3, ymax=1)
+
+    # 1111
+    bbox = binary_mask_segmentation._compute_bbox_from_rle(
+        rle_row_wise=[0, 4],
+        width=4,
+        height=1,
+    )
+    assert bbox == BoundingBox(xmin=0, ymin=0, xmax=3, ymax=0)