SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-13 04:25:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
SikongJueluo
2026-03-28 17:32:15 +08:00
parent 3c9a6f6eaf
commit f604c85a79
7 changed files with 252 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
mini-nav/compressors/pipeline.py
View File

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

19
mini-nav/scenegraph/roomnode.py
View File

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

13
mini-nav/simulator/topdown.py
View File

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

35
mini-nav/tests/test_image_utils.py
View File

@@ -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

3
mini-nav/utils/__init__.py
View File

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

22
mini-nav/utils/image.py
View File

@@ -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

87
notebooks/verification.py
View File

@@ -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,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,
semantic_hash=bits_np,
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__":