From cb93d83868192130129cca57b81a31a13e304dd6 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Mon, 30 Mar 2026 21:49:32 +0800 Subject: [PATCH] feat(simulator): add image saving utilities for verification --- mini-nav/simulator/__init__.py | 5 ++- mini-nav/simulator/image_save.py | 59 ++++++++++++++++++++++++++++++++ mini-nav/utils/__init__.py | 2 ++ mini-nav/utils/image.py | 21 ++++++++++++ notebooks/verification.py | 38 ++++++++++---------- 5 files changed, 105 insertions(+), 20 deletions(-) create mode 100644 mini-nav/simulator/image_save.py create mode 100644 mini-nav/utils/image.py diff --git a/mini-nav/simulator/__init__.py b/mini-nav/simulator/__init__.py index df74d74..a5e55e1 100644 --- a/mini-nav/simulator/__init__.py +++ b/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", ] diff --git a/mini-nav/simulator/image_save.py b/mini-nav/simulator/image_save.py new file mode 100644 index 0000000..159165f --- /dev/null +++ b/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 diff --git a/mini-nav/utils/__init__.py b/mini-nav/utils/__init__.py index 1591fce..b0ba98e 100644 --- a/mini-nav/utils/__init__.py +++ b/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", ] diff --git a/mini-nav/utils/image.py b/mini-nav/utils/image.py new file mode 100644 index 0000000..9c9ad7e --- /dev/null +++ b/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)) diff --git a/notebooks/verification.py b/notebooks/verification.py index 740d635..094e9a5 100644 --- a/notebooks/verification.py +++ b/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: