diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index 61d1e91..a53eec9 100644 --- a/mini-nav/compressors/pipeline.py +++ b/mini-nav/compressors/pipeline.py @@ -1,6 +1,6 @@ """SAM + DINO + Hash compression pipeline.""" -from typing import Optional +from typing import Optional, Sequence import torch import torch.nn as nn @@ -8,7 +8,7 @@ import torch.nn.functional as F from PIL import Image from utils import get_device -from utils.image import extract_masked_region, segment_image +from utils.image import extract_masked_region, segment_image, segment_image_dataset from utils.model import ( get_dino_dim, load_dino_model, @@ -105,6 +105,7 @@ class HashPipeline(nn.Module): image, min_area=self.sam_min_mask_area, max_masks=self.sam_max_masks, + points_per_batch=self.sam_points_per_batch, ) if not masks: @@ -112,6 +113,23 @@ class HashPipeline(nn.Module): return extract_masked_region(image, masks[0]["segment"]) + def _segment_with_sam_dataset( + self, + images: Sequence[Image.Image], + ) -> list[Image.Image]: + image_list = list(images) + masks_dataset = segment_image_dataset( + self.mask_generator, + image_list, + min_area=self.sam_min_mask_area, + max_masks=self.sam_max_masks, + points_per_batch=self.sam_points_per_batch, + ) + return [ + extract_masked_region(image, masks[0]["segment"]) if masks else image + for image, masks in zip(image_list, masks_dataset) + ] + def _dino_forward(self, image: Image.Image) -> torch.Tensor: """Extract DINO tokens from an image. @@ -127,6 +145,15 @@ class HashPipeline(nn.Module): outputs = self.dino(**inputs) return outputs.last_hidden_state + def _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor: + inputs = self.processor(images=list(images), return_tensors="pt").to( + self.device + ) + + with torch.no_grad(): + outputs = self.dino(**inputs) + return outputs.last_hidden_state + def forward(self, image: Image.Image) -> torch.Tensor: """Process a single image through the full pipeline. @@ -141,6 +168,36 @@ class HashPipeline(nn.Module): _, _, bits = self.hash_compressor(tokens) return bits + def forward_dataset( + self, + images: Sequence[Image.Image], + batch_size: int = 32, + apply_sam: bool = True, + ) -> torch.Tensor: + if batch_size <= 0: + raise ValueError("batch_size must be greater than 0") + + image_list = list(images) + + if len(image_list) == 0: + return torch.empty( + (0, self.hash_bits), dtype=torch.int32, device=self.device + ) + + if apply_sam: + processed_images = self._segment_with_sam_dataset(image_list) + else: + processed_images = image_list + + batch_bits: list[torch.Tensor] = [] + for i in range(0, len(processed_images), batch_size): + batch_images = processed_images[i : i + batch_size] + tokens = self._dino_forward_batch(batch_images) + _, _, bits = self.hash_compressor(tokens) + batch_bits.append(bits) + + return torch.cat(batch_bits, dim=0) + def extract_features(self, image: Image.Image) -> torch.Tensor: """Extract normalized DINO features from an image. @@ -154,3 +211,33 @@ class HashPipeline(nn.Module): tokens = self._dino_forward(image) features = tokens.mean(dim=1) return F.normalize(features, dim=-1) + + def extract_features_dataset( + self, + images: Sequence[Image.Image], + batch_size: int = 32, + apply_sam: bool = True, + ) -> torch.Tensor: + if batch_size <= 0: + raise ValueError("batch_size must be greater than 0") + + image_list = list(images) + + if len(image_list) == 0: + return torch.empty( + (0, self.dino_dim), dtype=torch.float32, device=self.device + ) + + if apply_sam: + processed_images = self._segment_with_sam_dataset(image_list) + else: + processed_images = image_list + + all_features: list[torch.Tensor] = [] + for i in range(0, len(processed_images), batch_size): + batch_images = processed_images[i : i + batch_size] + tokens = self._dino_forward_batch(batch_images) + features = tokens.mean(dim=1) + all_features.append(F.normalize(features, dim=-1)) + + return torch.cat(all_features, dim=0) diff --git a/mini-nav/scenegraph/roomnode.py b/mini-nav/scenegraph/roomnode.py index 702dc7e..9deb182 100644 --- a/mini-nav/scenegraph/roomnode.py +++ b/mini-nav/scenegraph/roomnode.py @@ -10,3 +10,22 @@ class RoomNode: # 范围：不用复杂的 Polygon，用一个中心点+半径，或者简单的 3D BBox 足够了 center: np.ndarray # [x, y, z] bbox_extent: np.ndarray # [dx, dy, dz] 用于快速判断一个点在不在房间里 + + def __post_init__(self): + self.center = np.asarray(self.center, dtype=np.float32) + self.bbox_extent = np.asarray(self.bbox_extent, dtype=np.float32) + + if self.center.shape != (3,): + raise ValueError(f"center must have shape (3,), got {self.center.shape}") + + if self.bbox_extent.shape != (3,): + raise ValueError(f"bbox_extent must have shape (3,), got {self.bbox_extent.shape}") + + if not np.all(np.isfinite(self.center)): + raise ValueError("center must contain only finite values") + + if not np.all(np.isfinite(self.bbox_extent)): + raise ValueError("bbox_extent must contain only finite values") + + if np.any(self.bbox_extent < 0): + raise ValueError("bbox_extent must be non-negative") \ No newline at end of file diff --git a/mini-nav/simulator/topdown.py b/mini-nav/simulator/topdown.py index 5dee7c5..9700ffd 100644 --- a/mini-nav/simulator/topdown.py +++ b/mini-nav/simulator/topdown.py @@ -30,6 +30,7 @@ def render_topdown_scene_map( style: TopDownRenderStyle | None = None, maps_module: Any | None = None, plt_module: Any | None = None, + use_matplotlib: bool = True, ) -> Any: if not elements.room_nodes: raise ValueError("room_nodes must not be empty") @@ -49,9 +50,6 @@ def render_topdown_scene_map( if maps_module is None: maps_module = import_module("habitat.utils.visualizations.maps") - if plt_module is None: - plt_module = import_module("matplotlib.pyplot") - map_height = float(elements.room_nodes[0].center[1]) top_down_map = maps_module.get_topdown_map( pathfinder, @@ -59,6 +57,15 @@ def render_topdown_scene_map( meters_per_pixel=meters_per_pixel, ) + if not use_matplotlib: + return top_down_map + + if plt_module is None: + try: + plt_module = import_module("matplotlib.pyplot") + except ImportError: + return top_down_map + plt_module.figure(figsize=style.figure_size) plt_module.imshow(top_down_map, cmap=style.map_cmap) diff --git a/mini-nav/tests/test_image_utils.py b/mini-nav/tests/test_image_utils.py index 6f7ec25..42439f9 100644 --- a/mini-nav/tests/test_image_utils.py +++ b/mini-nav/tests/test_image_utils.py @@ -3,7 +3,7 @@ from unittest.mock import Mock import torch from PIL import Image -from utils.image import segment_image +from utils.image import segment_image, segment_image_dataset def test_segment_image_passes_pil_image_to_mask_generator() -> None: @@ -80,3 +80,36 @@ def test_segment_image_filters_tensor_masks_by_min_area() -> None: assert len(result) == 1 assert result[0]["area"] == 4 + + +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]) + images = [ + Image.new("RGB", (3, 3), color=(0, 0, 0)), + Image.new("RGB", (3, 3), color=(0, 0, 0)), + ] + + result = segment_image_dataset( + mock_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 + assert mock_generator.call_count == 2 diff --git a/mini-nav/utils/__init__.py b/mini-nav/utils/__init__.py index 835da16..792a8de 100644 --- a/mini-nav/utils/__init__.py +++ b/mini-nav/utils/__init__.py @@ -4,7 +4,7 @@ from .feature_extractor import ( extract_single_image_feature, infer_vector_dim, ) -from .image import segment_image, extract_masked_region +from .image import extract_masked_region, segment_image, segment_image_dataset from .model import get_dino_dim, load_dino_model, load_hash_compressor, load_sam_model __all__ = [ @@ -14,6 +14,7 @@ __all__ = [ "extract_single_image_feature", "extract_batch_features", "segment_image", + "segment_image_dataset", "extract_masked_region", "load_dino_model", "load_sam_model", diff --git a/mini-nav/utils/image.py b/mini-nav/utils/image.py index 4bd61c3..94eea7a 100644 --- a/mini-nav/utils/image.py +++ b/mini-nav/utils/image.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Sequence import numpy as np from PIL import Image @@ -64,6 +64,26 @@ def segment_image( 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 eccbc17..e67c401 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -19,6 +19,8 @@ def import_packages(): import marimo as mo import numpy as np import polars as pl + from habitat.utils.visualizations import maps + from matplotlib import pyplot as plt from PIL import Image from compressors.pipeline import HashPipeline @@ -33,7 +35,6 @@ def import_packages(): from utils.image import extract_masked_region, segment_image return ( - BytesIO, HabitatSimulatorConfig, HashPipeline, Image, @@ -44,9 +45,11 @@ def import_packages(): collect_room_views_by_room, create_habitat_simulator, extract_masked_region, + maps, mo, np, pl, + plt, render_topdown_scene_map, segment_image, ) @@ -54,7 +57,10 @@ def import_packages(): @app.cell def setup_verification_context( - HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np + HabitatSimulatorConfig, + RoomNode, + create_habitat_simulator, + np, ): scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb" image_size = 256 @@ -65,6 +71,7 @@ def setup_verification_context( sam_max_masks = 5 sam_min_area = 32 * 32 hash_bits = 512 + pipeline_batch_size = 64 sim, agent = create_habitat_simulator( HabitatSimulatorConfig( @@ -91,11 +98,11 @@ def setup_verification_context( print("Sampled room centers:") for node in room_nodes: print(node.room_id, node.center) - return ( agent, hash_bits, meters_per_pixel, + pipeline_batch_size, room_nodes, sam_max_masks, sam_min_area, @@ -107,7 +114,9 @@ def setup_verification_context( @app.cell def render_topdown_room_map( TopDownSceneElements, + maps, meters_per_pixel, + plt, render_topdown_scene_map, room_nodes, sim, @@ -116,22 +125,25 @@ def render_topdown_room_map( pathfinder=sim.pathfinder, elements=TopDownSceneElements(room_nodes=room_nodes), meters_per_pixel=meters_per_pixel, + maps_module=maps, + plt_module=plt, ) return @app.cell def build_scene_graph_pipeline( - agent, HashPipeline, Image, ObjectNode, SimpleSceneGraph, + agent, collect_room_views_by_room, extract_masked_region, hash_bits, mo, np, + pipeline_batch_size, room_nodes, sam_max_masks, sam_min_area, @@ -161,16 +173,17 @@ def build_scene_graph_pipeline( total_masks = 0 object_index = 0 - view_jobs = [ + room_view_dataset = [ (room_id, view_idx, rgb) for room_id, views in all_room_views.items() for view_idx, rgb in enumerate(views) ] + object_dataset = [] for room_id, _view_idx, rgb in mo.status.progress_bar( - view_jobs, - title="Extracting masks and hashes", - subtitle="Running SAM + HashPipeline", + room_view_dataset, + title="Building object dataset", + subtitle="Running SAM segmentation", show_eta=True, show_rate=True, ): @@ -188,31 +201,56 @@ def build_scene_graph_pipeline( for mask in masks: masked_image = extract_masked_region(image, mask["segment"]) - bits = hash_pipeline(masked_image) + object_dataset.append((room_id, mask["bbox"], masked_image)) - bbox = mask["bbox"] - obj_center = np.array( - [bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2, 0.0], - dtype=np.float32, + if object_dataset: + masked_images = [item[2] for item in object_dataset] + batched_bits = hash_pipeline.forward_dataset( + masked_images, + batch_size=pipeline_batch_size, + apply_sam=False, + ) + if len(batched_bits) != len(object_dataset): + raise RuntimeError( + "Batch output size mismatch between masked images and hash outputs." + ) + else: + batched_bits = [] + + for ob_idx, (room_id, bbox, _) 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], + dtype=np.float32, + ) + + obj_id = f"obj_{object_index:04d}" + object_index += 1 + + bits_array = np.asarray(bits.detach().cpu().numpy()).reshape(-1) + if bits_array.size == 512: + bits_binary = (bits_array > 0).astype(np.uint8) + hash_bytes = np.packbits(bits_binary).tobytes() + elif bits_array.size == 64: + hash_bytes = bits_array.astype(np.uint8).tobytes() + else: + raise ValueError( + f"Unexpected hash length: {bits_array.size}. Expected 512 bits or 64 bytes." ) - obj_id = f"obj_{object_index:04d}" - object_index += 1 - bits_np = bits.squeeze().detach().cpu().numpy() - - scene_graph.objects[obj_id] = ObjectNode( - obj_id=obj_id, - room_id=room_id, - position=obj_center, - visual_hash=bits_np, - semantic_hash=bits_np, - hit_count=1, - last_seen_frame=0, - ) + scene_graph.objects[obj_id] = ObjectNode( + obj_id=obj_id, + room_id=room_id, + position=obj_center, + visual_hash=hash_bytes, + semantic_hash=hash_bytes, + hit_count=1, + last_seen_frame=0, + ) print(f"Total objects created: {len(scene_graph.objects)}") print(f"Total processed masks: {total_masks}") - return all_room_views, hash_pipeline, scene_graph + return (scene_graph,) @app.cell @@ -236,8 +274,8 @@ def build_room_and_object_tables(pl, scene_graph): "room_id": obj.room_id, "last_seen_frame": int(obj.last_seen_frame), "hit_count": int(obj.hit_count), - "visual_hash": obj.visual_hash.tolist(), - "semantic_hash": obj.semantic_hash.tolist(), + "visual_hash": obj.visual_hash.hex(), + "semantic_hash": obj.semantic_hash.hex(), } for obj in scene_graph.objects.values() ] @@ -248,22 +286,25 @@ def build_room_and_object_tables(pl, scene_graph): @app.cell -def upload_query_image(BytesIO, Image, mo, np): +def upload_query_image(mo): file_upload = mo.ui.file( filetypes=["image/*"], kind="area", label="Upload a query image", ) file_upload + return (file_upload,) - uploaded_image = None + +@app.cell +def _(file_upload, mo): + upload_image = None if file_upload.value: - contents = file_upload.contents() - if contents: - uploaded_image = Image.open(BytesIO(contents)) - mo.image(np.array(uploaded_image), alt="Uploaded query image") + upload_image = mo.image(file_upload.contents(), alt="Uploaded query image") - return file_upload, uploaded_image + # Build a grid. + upload_image if upload_image is not None else mo.md("No image uploaded yet.") + return if __name__ == "__main__":