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