Store semantic segmentation as run-length encoding. (#53)

Author: michal-lightly

.github/pull_request_template.md

@@ -0,0 +1,7 @@
+## What has changed and why?
+
+(Delete this: Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.)
+
+## How has it been tested?
+
+(Delete this: Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration.)

.python-version

@@ -1 +1 @@
-3.7.16
+3.8

src/labelformat/formats/semantic_segmentation/pascalvoc.py

@@ -114,7 +114,7 @@ def get_mask(self, image_filepath: str) -> SemanticSegmentationMask:
             valid_class_ids={c.id for c in self._categories},
         )
 
-        return SemanticSegmentationMask(array=mask_np)
+        return SemanticSegmentationMask.from_array(array=mask_np)
 
 
 def _validate_mask(

src/labelformat/model/semantic_segmentation.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import List, Optional, Tuple
+
 """Semantic segmentation core types and input interface.
 """
 
@@ -24,12 +26,36 @@ class SemanticSegmentationMask:
         array: The 2D numpy array with integer class IDs of shape (H, W).
     """
 
-    array: NDArray[np.int_]
+    category_id_rle: List[Tuple[int, int]]
+    """The mask as a run-length encoding (RLE) list of (category_id, run_length) tuples."""
+    width: int
+    height: int
 
-    def __post_init__(self) -> None:
-        if self.array.ndim != 2:
+    @classmethod
+    def from_array(cls, array: NDArray[np.int_]) -> "SemanticSegmentationMask":
+        """Create a SemanticSegmentationMask from a 2D numpy array."""
+        if array.ndim != 2:
             raise ValueError("SemSegMask.array must be 2D with shape (H, W).")
 
+        category_id_rle: List[Tuple[int, int]] = []
+
+        cur_cat_id: Optional[int] = None
+        cur_run_length = 0
+        for cat_id in array.flatten():
+            if cat_id == cur_cat_id:
+                cur_run_length += 1
+            else:
+                if cur_cat_id is not None:
+                    category_id_rle.append((cur_cat_id, cur_run_length))
+                cur_cat_id = cat_id
+                cur_run_length = 1
+        if cur_cat_id is not None:
+            category_id_rle.append((cur_cat_id, cur_run_length))
+
+        return cls(
+            category_id_rle=category_id_rle, width=array.shape[1], height=array.shape[0]
+        )
+
 
 class SemanticSegmentationInput(ABC):
 

tests/unit/formats/semantic_segmentation/test_pascalvoc.py

@@ -45,17 +45,16 @@ def test_from_dirs__builds_categories_and_images(self) -> None:
         filenames = {img.filename for img in imgs}
         assert filenames == {"2007_000032.jpg", "subdir/2007_000033.jpg"}
 
-    def test_get_mask__returns_int2d_and_matches_image_shape(self) -> None:
+    def test_get_mask__returns_rle_and_matches_image_length(self) -> None:
         mapping = _load_class_mapping_int_keys()
         ds = PascalVOCSemanticSegmentationInput.from_dirs(
             images_dir=IMAGES_DIR, masks_dir=MASKS_DIR, class_id_to_name=mapping
         )
 
         for img in ds.get_images():
             mask = ds.get_mask(img.filename)
-            assert mask.array.ndim == 2
-            assert np.issubdtype(mask.array.dtype, np.integer)
-            assert mask.array.shape == (img.height, img.width)
+            length = sum(run_length for _, run_length in mask.category_id_rle)
+            assert length == img.width * img.height
 
     def test_from_dirs__missing_mask_raises(self, tmp_path: Path) -> None:
         masks_tmp = tmp_path / "SegmentationClass"

tests/unit/model/test_semantic_segmentation.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import numpy as np
+
+from labelformat.model.semantic_segmentation import SemanticSegmentationMask
+
+
+class TestSemanticSegmentationMask:
+    def test_from_array(self) -> None:
+        array = np.array(
+            [
+                [1, 1, 2, 2],
+                [2, 1, 1, 1],
+                [3, 3, 3, 3],
+            ],
+            dtype=np.int_,
+        )
+        expected_rle = [
+            (1, 2),
+            (2, 3),
+            (1, 3),
+            (3, 4),
+        ]
+        mask = SemanticSegmentationMask.from_array(array=array)
+        assert mask.category_id_rle == expected_rle
+        assert mask.width == 4
+        assert mask.height == 3