feat(verification): add batch segmentation and image saving

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,)