feat(verification): add batch segmentation and image saving

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