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

notebooks/verification.py

@@ -23,6 +23,7 @@ def import_packages():
     from matplotlib import pyplot as plt
     from PIL import Image
 
+    from configs import cfg_manager
     from compressors.pipeline import HashPipeline
     from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
     from simulator import (
@@ -32,7 +33,7 @@ def import_packages():
         create_habitat_simulator,
         render_topdown_scene_map,
     )
-    from utils.image import extract_masked_region, segment_image
+    from utils.image import extract_masked_region, segment_image_dataset
 
     return (
         HabitatSimulatorConfig,
@@ -44,6 +45,7 @@ def import_packages():
         TopDownSceneElements,
         collect_room_views_by_room,
         create_habitat_simulator,
+        cfg_manager,
         extract_masked_region,
         maps,
         mo,
@@ -51,7 +53,7 @@ def import_packages():
         pl,
         plt,
         render_topdown_scene_map,
-        segment_image,
+        segment_image_dataset,
     )
 
 
@@ -139,6 +141,7 @@ def build_scene_graph_pipeline(
     SimpleSceneGraph,
     agent,
     collect_room_views_by_room,
+    cfg_manager,
     extract_masked_region,
     hash_bits,
     mo,
@@ -147,7 +150,7 @@ def build_scene_graph_pipeline(
     room_nodes,
     sam_max_masks,
     sam_min_area,
-    segment_image,
+    segment_image_dataset,
     sim,
     views_per_room,
 ):
@@ -170,6 +173,9 @@ def build_scene_graph_pipeline(
         rooms={room.room_id: room for room in room_nodes},
         objects={},
     )
+    verification_output_dir = cfg_manager.get().output.directory / "verification"
+    verification_output_dir.mkdir(parents=True, exist_ok=True)
+
     total_masks = 0
     object_index = 0
 
@@ -180,31 +186,48 @@ def build_scene_graph_pipeline(
     ]
     object_dataset = []
 
-    for room_id, _view_idx, rgb in mo.status.progress_bar(
-        room_view_dataset,
+    room_view_images = []
+    for _, _, rgb in room_view_dataset:
+        rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
+        room_view_images.append(Image.fromarray(rgb3.astype(np.uint8)))
+
+    masks_dataset = segment_image_dataset(
+        hash_pipeline.mask_generator,
+        room_view_images,
+        min_area=hash_pipeline.sam_min_mask_area,
+        max_masks=hash_pipeline.sam_max_masks,
+        points_per_batch=hash_pipeline.sam_points_per_batch,
+    )
+    if len(masks_dataset) != len(room_view_dataset):
+        raise RuntimeError("SAM dataset output size mismatch with room_view_dataset.")
+
+    dataset_jobs = list(zip(room_view_dataset, room_view_images, masks_dataset))
+    for (room_id, view_idx, _), image, masks in mo.status.progress_bar(
+        dataset_jobs,
         title="Building object dataset",
         subtitle="Running SAM segmentation",
         show_eta=True,
         show_rate=True,
     ):
-        rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
-        image = Image.fromarray(rgb3.astype(np.uint8))
+        room_output_dir = verification_output_dir / room_id
+        room_output_dir.mkdir(parents=True, exist_ok=True)
+        room_view_path = room_output_dir / f"view_{view_idx:03d}.png"
+        image.convert("RGB").save(room_view_path, format="PNG")
 
-        masks = segment_image(
-            hash_pipeline.mask_generator,
-            image,
-            min_area=hash_pipeline.sam_min_mask_area,
-            max_masks=hash_pipeline.sam_max_masks,
-            points_per_batch=hash_pipeline.sam_points_per_batch,
-        )
         total_masks += len(masks)
 
-        for mask in masks:
+        for mask_idx, mask in enumerate(masks):
             masked_image = extract_masked_region(image, mask["segment"])
-            object_dataset.append((room_id, mask["bbox"], masked_image))
+            object_dataset.append(
+                (room_id, view_idx, mask_idx, mask["bbox"], masked_image)
+            )
 
     if object_dataset:
-        masked_images = [item[2] for item in object_dataset]
+        masked_images = [item[4] for item in object_dataset]
+        if any(not isinstance(img, Image.Image) for img in masked_images):
+            raise TypeError(
+                "object_dataset contains non-image entries for batch inference."
+            )
         batched_bits = hash_pipeline.forward_dataset(
             masked_images,
             batch_size=pipeline_batch_size,
@@ -217,7 +240,9 @@ def build_scene_graph_pipeline(
     else:
         batched_bits = []
 
-    for ob_idx, (room_id, bbox, _) in enumerate(object_dataset):
+    for ob_idx, (room_id, view_idx, mask_idx, bbox, masked_image) in enumerate(
+        object_dataset
+    ):
         bits = batched_bits[ob_idx]
         obj_center = np.array(
             [bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2, 0.0],
@@ -227,6 +252,13 @@ def build_scene_graph_pipeline(
         obj_id = f"obj_{object_index:04d}"
         object_index += 1
 
+        room_output_dir = verification_output_dir / room_id
+        room_output_dir.mkdir(parents=True, exist_ok=True)
+        object_image_path = (
+            room_output_dir / f"{obj_id}_view{view_idx:03d}_mask{mask_idx:02d}.png"
+        )
+        masked_image.convert("RGB").save(object_image_path, format="PNG")
+
         bits_array = np.asarray(bits.detach().cpu().numpy()).reshape(-1)
         if bits_array.size == 512:
             bits_binary = (bits_array > 0).astype(np.uint8)
@@ -250,6 +282,7 @@ def build_scene_graph_pipeline(
 
     print(f"Total objects created: {len(scene_graph.objects)}")
     print(f"Total processed masks: {total_masks}")
+    print(f"Saved object images to: {verification_output_dir}")
     return (scene_graph,)