feat(pipeline): add batch processing for scene graph construction

mini-nav/tests/test_image_utils.py

@@ -3,7 +3,7 @@ from unittest.mock import Mock
 import torch
 from PIL import Image
 
-from utils.image import segment_image
+from utils.image import segment_image, segment_image_dataset
 
 
 def test_segment_image_passes_pil_image_to_mask_generator() -> None:
@@ -80,3 +80,36 @@ def test_segment_image_filters_tensor_masks_by_min_area() -> None:
 
     assert len(result) == 1
     assert result[0]["area"] == 4
+
+
+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])
+    images = [
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+    ]
+
+    result = segment_image_dataset(
+        mock_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
+    assert mock_generator.call_count == 2