From f6c1a67e8831ef798bd65d9032687260cb20a313 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Sat, 28 Mar 2026 21:30:02 +0800 Subject: [PATCH] feat(verification): add batch segmentation and image saving --- mini-nav/simulator/habitat.py | 2 +- mini-nav/tests/test_image_utils.py | 69 ++++++++++++++++---- mini-nav/utils/image.py | 100 +++++++++++++++++++++-------- notebooks/verification.py | 69 ++++++++++++++------ 4 files changed, 182 insertions(+), 58 deletions(-) diff --git a/mini-nav/simulator/habitat.py b/mini-nav/simulator/habitat.py index dcf4e11..ec15fce 100644 --- a/mini-nav/simulator/habitat.py +++ b/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 diff --git a/mini-nav/tests/test_image_utils.py b/mini-nav/tests/test_image_utils.py index 42439f9..8e3093d 100644 --- a/mini-nav/tests/test_image_utils.py +++ b/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 diff --git a/mini-nav/utils/image.py b/mini-nav/utils/image.py index 94eea7a..d6ed081 100644 --- a/mini-nav/utils/image.py +++ b/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 diff --git a/notebooks/verification.py b/notebooks/verification.py index e67c401..b86f810 100644 --- a/notebooks/verification.py +++ b/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,)