feat(verification): add batch segmentation and image saving

mini-nav/simulator/habitat.py

@@ -9,7 +9,7 @@ from typing import Any
 class HabitatSimulatorConfig:
     scene_path: str
     views_per_room: int = 6
-    image_size: int = 256
+    image_size: int = 512
     sensor_height: float = 1.5
     move_forward_step: float = 0.25
     enable_physics: bool = False

mini-nav/tests/test_image_utils.py

@@ -83,19 +83,17 @@ def test_segment_image_filters_tensor_masks_by_min_area() -> None:
 
 
 def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
-    first_masks = {
-        "masks": torch.tensor(
-            [[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
-            dtype=torch.float32,
-        )
-    }
-    second_masks = {
-        "masks": torch.tensor(
-            [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
-            dtype=torch.float32,
-        )
-    }
-    mock_generator = Mock(side_effect=[first_masks, second_masks])
+    first_masks = torch.tensor(
+        [[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
+        dtype=torch.float32,
+    )
+    second_masks = torch.tensor(
+        [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
+        dtype=torch.float32,
+    )
+    mock_generator = Mock(
+        return_value=[{"masks": first_masks}, {"masks": second_masks}]
+    )
     images = [
         Image.new("RGB", (3, 3), color=(0, 0, 0)),
         Image.new("RGB", (3, 3), color=(0, 0, 0)),
@@ -112,4 +110,47 @@ def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
     assert len(result) == 2
     assert result[0][0]["area"] == 4
     assert result[1][0]["area"] == 9
-    assert mock_generator.call_count == 2
+    assert mock_generator.call_count == 1
+
+
+def test_segment_image_dataset_falls_back_to_single_image_calls() -> None:
+    call_index = {"value": 0}
+
+    def fake_generator(images, points_per_batch):
+        if isinstance(images, list):
+            raise TypeError("Batch input unsupported")
+
+        result_options = [
+            {
+                "masks": torch.tensor(
+                    [[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
+                    dtype=torch.float32,
+                )
+            },
+            {
+                "masks": torch.tensor(
+                    [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
+                    dtype=torch.float32,
+                )
+            },
+        ]
+        out = result_options[call_index["value"]]
+        call_index["value"] += 1
+        return out
+
+    images = [
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+    ]
+
+    result = segment_image_dataset(
+        fake_generator,
+        images,
+        min_area=2,
+        max_masks=5,
+        points_per_batch=16,
+    )
+
+    assert len(result) == 2
+    assert result[0][0]["area"] == 4
+    assert result[1][0]["area"] == 9

mini-nav/utils/image.py

@@ -29,10 +29,79 @@ def segment_image(
     """
     image_rgb = image.convert("RGB")
     raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch)
-    raw_masks = raw_output.get("masks", raw_output)
+    return _normalize_and_filter_masks(
+        raw_output, min_area=min_area, max_masks=max_masks
+    )
+
+
+def segment_image_dataset(
+    mask_generator: Any,
+    images: Sequence[Image.Image],
+    min_area: int = 32 * 32,
+    max_masks: int = 5,
+    points_per_batch: int = 64,
+) -> list[list[dict[str, Any]]]:
+    image_list = list(images)
+    if not image_list:
+        return []
+
+    image_rgb_list = [image.convert("RGB") for image in image_list]
+    try:
+        raw_batch_output = mask_generator(
+            image_rgb_list,
+            points_per_batch=points_per_batch,
+        )
+        batch_items = _split_batch_output(
+            raw_batch_output, expected_size=len(image_list)
+        )
+        if batch_items is not None:
+            return [
+                _normalize_and_filter_masks(
+                    batch_item,
+                    min_area=min_area,
+                    max_masks=max_masks,
+                )
+                for batch_item in batch_items
+            ]
+    except TypeError:
+        pass
+
+    return [
+        _normalize_and_filter_masks(
+            mask_generator(image_rgb, points_per_batch=points_per_batch),
+            min_area=min_area,
+            max_masks=max_masks,
+        )
+        for image_rgb in image_rgb_list
+    ]
+
+
+def _split_batch_output(raw_output: Any, expected_size: int) -> list[Any] | None:
+    if isinstance(raw_output, list):
+        if len(raw_output) == expected_size:
+            return raw_output
+        return None
+
+    if isinstance(raw_output, dict):
+        raw_masks = raw_output.get("masks", raw_output)
+        if isinstance(raw_masks, list) and len(raw_masks) == expected_size:
+            return raw_masks
+
+    return None
+
+
+def _normalize_and_filter_masks(
+    raw_output: Any,
+    min_area: int,
+    max_masks: int,
+) -> list[dict[str, Any]]:
+    raw_masks = (
+        raw_output.get("masks", raw_output)
+        if isinstance(raw_output, dict)
+        else raw_output
+    )
 
     normalized_masks: list[dict[str, Any]] = []
-
     if isinstance(raw_masks, list):
         if raw_masks and isinstance(raw_masks[0], dict):
             normalized_masks = raw_masks
@@ -55,35 +124,16 @@ def segment_image(
     if not normalized_masks:
         return []
 
-    filtered_masks = [m for m in normalized_masks if int(m["area"]) >= min_area]
-
+    filtered_masks = [
+        mask for mask in normalized_masks if int(mask["area"]) >= min_area
+    ]
     if not filtered_masks:
         return []
 
-    sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
+    sorted_masks = sorted(filtered_masks, key=lambda mask: mask["area"], reverse=True)
     return sorted_masks[:max_masks]
 
 
-def segment_image_dataset(
-    mask_generator: Any,
-    images: Sequence[Image.Image],
-    min_area: int = 32 * 32,
-    max_masks: int = 5,
-    points_per_batch: int = 64,
-) -> list[list[dict[str, Any]]]:
-    image_list = list(images)
-    return [
-        segment_image(
-            mask_generator,
-            image,
-            min_area=min_area,
-            max_masks=max_masks,
-            points_per_batch=points_per_batch,
-        )
-        for image in image_list
-    ]
-
-
 def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
     if mask_like is None:
         return None