feat(pipeline): add batch processing for scene graph construction

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)

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")

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)
 

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

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",

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

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,
+        room_view_dataset,
-        title="Extracting masks and hashes",
+        title="Building object dataset",
-        subtitle="Running SAM + HashPipeline",
+        subtitle="Running SAM segmentation",
         show_eta=True,
         show_rate=True,
     ):
@@ -188,9 +201,24 @@ 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"]
+    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,
@@ -198,21 +226,31 @@ def build_scene_graph_pipeline(
 
         obj_id = f"obj_{object_index:04d}"
         object_index += 1
-            bits_np = bits.squeeze().detach().cpu().numpy()
+
+        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."
+            )
 
         scene_graph.objects[obj_id] = ObjectNode(
             obj_id=obj_id,
             room_id=room_id,
             position=obj_center,
-                visual_hash=bits_np,
+            visual_hash=hash_bytes,
-                semantic_hash=bits_np,
+            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(),
+            "visual_hash": obj.visual_hash.hex(),
-            "semantic_hash": obj.semantic_hash.tolist(),
+            "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()
+        upload_image = mo.image(file_upload.contents(), alt="Uploaded query image")
-        if contents:
-            uploaded_image = Image.open(BytesIO(contents))
-            mo.image(np.array(uploaded_image), 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__":