feat(simulator): add image saving utilities for verification

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 = []
-    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)))
+    room_view_images = [numpy_to_pil(rgb) for _, _, rgb in room_view_dataset]
 
     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",
-        subtitle="Running SAM segmentation",
-        show_eta=True,
-        show_rate=True,
+        description="Building object dataset...",
     ):
-        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")
-
+        save_room_view(verification_output_dir, room_id, view_idx, image)
         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
-        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"
+        save_object_image(
+            verification_output_dir, room_id, obj_id, view_idx, mask_idx, masked_image
         )
-        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: