feat(simulator): add image saving utilities for verification

mini-nav/simulator/__init__.py

@@ -3,16 +3,19 @@ from .habitat import (
     close_habitat_simulator,
     create_habitat_simulator,
 )
+from .image_save import save_object_image, save_room_view
 from .topdown import TopDownRenderStyle, TopDownSceneElements, render_topdown_scene_map
 from .views import RoomViewsByRoom, collect_room_views_by_room
 
 __all__ = [
     "HabitatSimulatorConfig",
+    "RoomViewsByRoom",
     "TopDownRenderStyle",
     "TopDownSceneElements",
-    "RoomViewsByRoom",
     "close_habitat_simulator",
     "collect_room_views_by_room",
     "create_habitat_simulator",
     "render_topdown_scene_map",
+    "save_object_image",
+    "save_room_view",
 ]

mini-nav/simulator/image_save.py

@@ -0,0 +1,59 @@
+"""Image saving utilities for simulation verification output."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from PIL import Image
+
+
+def save_room_view(
+    output_dir: Path,
+    room_id: str,
+    view_idx: int,
+    image: Image.Image,
+) -> Path:
+    """Save a single room view image to disk.
+
+    Args:
+        output_dir: Base output directory.
+        room_id: Room identifier string.
+        view_idx: Zero-based view index within the room.
+        image: PIL Image to save.
+
+    Returns:
+        Path where the image was saved.
+    """
+    room_dir = output_dir / room_id
+    room_dir.mkdir(parents=True, exist_ok=True)
+    path = room_dir / f"view_{view_idx:03d}.png"
+    image.convert("RGB").save(path, format="PNG")
+    return path
+
+
+def save_object_image(
+    output_dir: Path,
+    room_id: str,
+    obj_id: str,
+    view_idx: int,
+    mask_idx: int,
+    image: Image.Image,
+) -> Path:
+    """Save a masked object image to disk.
+
+    Args:
+        output_dir: Base output directory.
+        room_id: Room identifier string.
+        obj_id: Object identifier string.
+        view_idx: Zero-based view index.
+        mask_idx: Zero-based mask index within the view.
+        image: PIL Image to save.
+
+    Returns:
+        Path where the image was saved.
+    """
+    room_dir = output_dir / room_id
+    room_dir.mkdir(parents=True, exist_ok=True)
+    path = room_dir / f"{obj_id}_view{view_idx:03d}_mask{mask_idx:02d}.png"
+    image.convert("RGB").save(path, format="PNG")
+    return path

mini-nav/utils/__init__.py

@@ -4,6 +4,7 @@ from .feature_extractor import (
     extract_single_image_feature,
     infer_vector_dim,
 )
+from .image import numpy_to_pil
 
 __all__ = [
     "get_device",
@@ -11,4 +12,5 @@ __all__ = [
     "infer_vector_dim",
     "extract_single_image_feature",
     "extract_batch_features",
+    "numpy_to_pil",
 ]

mini-nav/utils/image.py

@@ -0,0 +1,21 @@
+"""Image conversion utilities."""
+
+from __future__ import annotations
+
+import numpy as np
+from PIL import Image
+
+
+def numpy_to_pil(rgb: np.ndarray) -> Image.Image:
+    """Convert an RGB numpy array to a PIL Image.
+
+    Handles arrays with 4 channels (RGBA) by dropping the alpha channel.
+
+    Args:
+        rgb: Numpy array of shape (H, W, C) with dtype uint8 or compatible.
+
+    Returns:
+        PIL Image in RGB mode.
+    """
+    rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
+    return Image.fromarray(rgb3.astype(np.uint8))

notebooks/verification.py

@@ -23,6 +23,8 @@ def import_packages():
     from matplotlib import pyplot as plt
     from PIL import Image
 
+    from rich.progress import track
+
     from configs import cfg_manager
     from compressors.pipeline import HashPipeline
     from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
@@ -32,8 +34,11 @@ def import_packages():
         collect_room_views_by_room,
         create_habitat_simulator,
         render_topdown_scene_map,
+        save_object_image,
+        save_room_view,
     )
     from compressors.proposal import extract_masked_region, generate_proposals_batch
+    from utils import numpy_to_pil
 
     return (
         HabitatSimulatorConfig,
@@ -50,10 +55,14 @@ def import_packages():
         maps,
         mo,
         np,
+        numpy_to_pil,
         pl,
         plt,
         render_topdown_scene_map,
+        save_object_image,
+        save_room_view,
         generate_proposals_batch,
+        track,
     )
 
 
@@ -146,12 +155,16 @@ def build_scene_graph_pipeline(
     hash_bits,
     mo,
     np,
+    numpy_to_pil,
     pipeline_batch_size,
     room_nodes,
     sam_max_masks,
     sam_min_area,
     generate_proposals_batch,
+    save_object_image,
+    save_room_view,
     sim,
+    track,
     views_per_room,
 ):
     all_room_views = collect_room_views_by_room(
@@ -186,10 +199,7 @@ def build_scene_graph_pipeline(
     ]
     object_dataset = []
 
-    room_view_images = []
+    room_view_images = [numpy_to_pil(rgb) for _, _, rgb in room_view_dataset]
-    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 = generate_proposals_batch(
         hash_pipeline.mask_generator,
@@ -202,18 +212,11 @@ def build_scene_graph_pipeline(
         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(
+    for (room_id, view_idx, _), image, masks in track(
         dataset_jobs,
-        title="Building object dataset",
+        description="Building object dataset...",
-        subtitle="Running SAM segmentation",
-        show_eta=True,
-        show_rate=True,
     ):
-        room_output_dir = verification_output_dir / room_id
+        save_room_view(verification_output_dir, room_id, view_idx, image)
-        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")
-
         total_masks += len(masks)
 
         for mask_idx, mask in enumerate(masks):
@@ -252,12 +255,9 @@ def build_scene_graph_pipeline(
         obj_id = f"obj_{object_index:04d}"
         object_index += 1
 
-        room_output_dir = verification_output_dir / room_id
+        save_object_image(
-        room_output_dir.mkdir(parents=True, exist_ok=True)
+            verification_output_dir, room_id, obj_id, view_idx, mask_idx, masked_image
-        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: