refactor(compressors): consolidate pipeline and improve mask handling

mini-nav/tests/test_image_utils.py

@@ -0,0 +1,82 @@
+from unittest.mock import Mock
+
+import torch
+from PIL import Image
+
+from utils.image import segment_image
+
+
+def test_segment_image_passes_pil_image_to_mask_generator() -> None:
+    mock_generator = Mock(return_value={"masks": []})
+
+    segment_image(
+        mock_generator,
+        Image.new("RGBA", (16, 16), color=(255, 0, 0, 255)),
+        points_per_batch=32,
+    )
+
+    image_arg = mock_generator.call_args.args[0]
+    assert isinstance(image_arg, Image.Image)
+    assert image_arg.mode == "RGB"
+    assert mock_generator.call_args.kwargs["points_per_batch"] == 32
+
+
+def test_segment_image_supports_tensor_masks_output() -> None:
+    masks_tensor = torch.tensor(
+        [
+            [
+                [1, 1, 0],
+                [1, 1, 0],
+                [0, 0, 0],
+            ],
+            [
+                [1, 1, 1],
+                [1, 1, 1],
+                [1, 1, 1],
+            ],
+        ],
+        dtype=torch.float32,
+    )
+    mock_generator = Mock(return_value={"masks": masks_tensor})
+
+    result = segment_image(
+        mock_generator,
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+        min_area=3,
+        max_masks=5,
+    )
+
+    assert len(result) == 2
+    assert result[0]["area"] == 9
+    assert result[0]["bbox"] == [0, 0, 3, 3]
+    assert result[1]["area"] == 4
+    assert result[1]["bbox"] == [0, 0, 2, 2]
+
+
+def test_segment_image_filters_tensor_masks_by_min_area() -> None:
+    masks_tensor = torch.tensor(
+        [
+            [
+                [1, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ],
+            [
+                [1, 1, 0],
+                [1, 1, 0],
+                [0, 0, 0],
+            ],
+        ],
+        dtype=torch.float32,
+    )
+    mock_generator = Mock(return_value={"masks": masks_tensor})
+
+    result = segment_image(
+        mock_generator,
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+        min_area=2,
+        max_masks=5,
+    )
+
+    assert len(result) == 1
+    assert result[0]["area"] == 4