mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
a127032e18
Introduce SceneGraphBuilder + SceneGraphBuildConfig to decouple scene graph construction from the verification notebook. The builder handles batch inference, hash encoding, and object node creation internally. - Add SceneGraphBuilder.build_from_room_views() as the main entry point - Add SceneGraphBuildConfig for inference_batch_size and position strategy - Add SceneGraphBuildArtifacts to carry cropped images and debug metadata - Extend ObjectNode with optional detection metadata (label, confidence, bbox_xyxy, source_view_id, position_confidence) - Add RoomView frozen dataclass as a structured view container - Add flatten_room_views() utility to replace inline list comprehensions - Refactor notebooks/verification.py to use the new builder API BREAKING CHANGE: ObjectNode now accepts additional optional fields; direct scene_graph.objects[obj_id] = ObjectNode(...) construction in the notebook is replaced by builder.build_from_room_views(...).
294 lines
11 KiB
Python
294 lines
11 KiB
Python
"""Scene graph builder: converts room views + pipeline debug output into SimpleSceneGraph."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Literal, Sequence
|
|
|
|
import numpy as np
|
|
|
|
from .hash_codec import bits_tensor_to_hash_bytes
|
|
from .objectnode import ObjectNode
|
|
from .roomnode import RoomNode
|
|
from .scenegraph import SimpleSceneGraph
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SceneGraphBuildConfig:
|
|
"""Configuration for the scene graph builder.
|
|
|
|
Attributes:
|
|
position_strategy: How object positions are assigned. Only 'room_center'
|
|
is supported in M0.
|
|
inference_batch_size: Batch size passed to the pipeline.
|
|
enable_fusion: Whether to fuse overlapping detections (not yet supported).
|
|
fusion_hash_similarity_threshold: Hash similarity threshold for fusion.
|
|
fusion_distance_threshold_m: Distance threshold in meters for fusion.
|
|
"""
|
|
|
|
position_strategy: Literal["room_center"] = "room_center"
|
|
inference_batch_size: int = 4
|
|
enable_fusion: bool = False
|
|
fusion_hash_similarity_threshold: float = 0.95
|
|
fusion_distance_threshold_m: float = 0.5
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.position_strategy != "room_center":
|
|
raise ValueError(
|
|
f"position_strategy must be 'room_center', got {self.position_strategy!r}"
|
|
)
|
|
if self.inference_batch_size <= 0:
|
|
raise ValueError(
|
|
f"inference_batch_size must be positive, got {self.inference_batch_size}"
|
|
)
|
|
if self.enable_fusion:
|
|
raise ValueError("fusion is not yet supported in M0")
|
|
|
|
|
|
@dataclass
|
|
class SceneGraphBuildArtifacts:
|
|
"""Artifacts produced during scene graph construction.
|
|
|
|
Attributes:
|
|
object_images: Maps object_id -> cropped image for each detected object.
|
|
debug_meta: Pipeline debug_metadata for each input view.
|
|
cropped_images: Flat list of all cropped images from the pipeline.
|
|
"""
|
|
|
|
object_images: dict[str, Any] = field(default_factory=dict)
|
|
debug_meta: list[dict[str, Any]] = field(default_factory=list)
|
|
cropped_images: list[Any] = field(default_factory=list)
|
|
|
|
|
|
class SceneGraphBuilder:
|
|
"""Builds a SimpleSceneGraph from room views and a detection pipeline."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
pipeline: Any,
|
|
config: SceneGraphBuildConfig | None = None,
|
|
) -> None:
|
|
self._pipeline = pipeline
|
|
self._config = config or SceneGraphBuildConfig()
|
|
|
|
def build_from_room_views(
|
|
self,
|
|
*,
|
|
room_nodes: Sequence[RoomNode],
|
|
room_views: Sequence[Any],
|
|
text_labels: list[str],
|
|
) -> tuple[SimpleSceneGraph, SceneGraphBuildArtifacts]:
|
|
"""Build a scene graph from room views and a list of room nodes.
|
|
|
|
Args:
|
|
room_nodes: Sequence of RoomNode objects describing each room.
|
|
room_views: Sequence of RoomView objects (must have .rgb, .room_id,
|
|
.view_idx attributes).
|
|
text_labels: Text labels to pass to the pipeline.
|
|
|
|
Returns:
|
|
A tuple of (SimpleSceneGraph, SceneGraphBuildArtifacts).
|
|
|
|
Raises:
|
|
ValueError: If any room_view references an unknown room_id, or if the
|
|
pipeline output is inconsistent.
|
|
"""
|
|
graph = self._prepare_graph(room_nodes=room_nodes, room_views=room_views)
|
|
output = self._run_pipeline(room_views=room_views, text_labels=text_labels)
|
|
artifacts = self._build_artifacts(output)
|
|
prefix = self._validate_pipeline_output(output=output, room_views=room_views)
|
|
self._add_objects_to_graph(
|
|
graph=graph,
|
|
artifacts=artifacts,
|
|
room_views=room_views,
|
|
output=output,
|
|
prefix=prefix,
|
|
)
|
|
return graph, artifacts
|
|
|
|
def _prepare_graph(
|
|
self,
|
|
*,
|
|
room_nodes: Sequence[RoomNode],
|
|
room_views: Sequence[Any],
|
|
) -> SimpleSceneGraph:
|
|
rooms = {node.room_id: node for node in room_nodes}
|
|
for view in room_views:
|
|
if view.room_id not in rooms:
|
|
raise ValueError(
|
|
f"Missing/unknown room {view.room_id!r} in room_view"
|
|
)
|
|
return SimpleSceneGraph(rooms=rooms, objects={})
|
|
|
|
def _run_pipeline(self, *, room_views: Sequence[Any], text_labels: list[str]) -> Any:
|
|
images = [view.rgb for view in room_views]
|
|
return self._pipeline.process_batch(
|
|
images,
|
|
text_labels,
|
|
batch_size=self._config.inference_batch_size,
|
|
return_debug_details=True,
|
|
)
|
|
|
|
def _build_artifacts(self, output: Any) -> SceneGraphBuildArtifacts:
|
|
return SceneGraphBuildArtifacts(
|
|
debug_meta=list(output.debug_meta),
|
|
cropped_images=list(output.cropped_images),
|
|
)
|
|
|
|
def _validate_pipeline_output(
|
|
self,
|
|
*,
|
|
output: Any,
|
|
room_views: Sequence[Any],
|
|
) -> list[int]:
|
|
self._validate_debug_meta_length(output=output, room_views=room_views)
|
|
num_selected_list = self._validate_selected_counts(output.debug_meta)
|
|
total_selected = sum(num_selected_list)
|
|
self._validate_crop_and_hash_counts(output=output, total_selected=total_selected)
|
|
return self._prefix_offsets(num_selected_list)
|
|
|
|
def _validate_debug_meta_length(
|
|
self,
|
|
*,
|
|
output: Any,
|
|
room_views: Sequence[Any],
|
|
) -> None:
|
|
if len(output.debug_meta) != len(room_views):
|
|
raise ValueError(
|
|
f"debug_meta length ({len(output.debug_meta)}) does not match "
|
|
f"room_views length ({len(room_views)})"
|
|
)
|
|
|
|
def _validate_selected_counts(self, debug_meta: Sequence[dict[str, Any]]) -> list[int]:
|
|
num_selected_list: list[int] = []
|
|
for view_idx, meta in enumerate(debug_meta):
|
|
selected_indices = meta.get("selected_indices", [])
|
|
num_selected = int(meta.get("num_selected", 0))
|
|
if len(selected_indices) != num_selected:
|
|
raise ValueError(
|
|
f"View {view_idx}: len(selected_indices) ({len(selected_indices)}) "
|
|
f"does not match num_selected ({num_selected})"
|
|
)
|
|
num_selected_list.append(num_selected)
|
|
return num_selected_list
|
|
|
|
def _validate_crop_and_hash_counts(self, *, output: Any, total_selected: int) -> None:
|
|
if total_selected != len(output.cropped_images):
|
|
raise ValueError(
|
|
f"total_selected ({total_selected}) does not match "
|
|
f"len(cropped_images) ({len(output.cropped_images)})"
|
|
)
|
|
|
|
hash_bits = output.hash_bits
|
|
if hash_bits is None:
|
|
if total_selected > 0:
|
|
raise ValueError(
|
|
f"hash_bits is None but total_selected ({total_selected}) > 0"
|
|
)
|
|
elif total_selected != hash_bits.shape[0]:
|
|
raise ValueError(
|
|
f"total_selected ({total_selected}) does not match "
|
|
f"hash_bits.shape[0] ({hash_bits.shape[0]})"
|
|
)
|
|
|
|
def _prefix_offsets(self, counts: Sequence[int]) -> list[int]:
|
|
prefix: list[int] = []
|
|
running = 0
|
|
for count in counts:
|
|
prefix.append(running)
|
|
running += count
|
|
return prefix
|
|
|
|
def _add_objects_to_graph(
|
|
self,
|
|
*,
|
|
graph: SimpleSceneGraph,
|
|
artifacts: SceneGraphBuildArtifacts,
|
|
room_views: Sequence[Any],
|
|
output: Any,
|
|
prefix: Sequence[int],
|
|
) -> None:
|
|
for image_idx, (view, meta) in enumerate(zip(room_views, output.debug_meta)):
|
|
for local_mask_idx, selected_idx in enumerate(meta.get("selected_indices", [])):
|
|
crop_index = prefix[image_idx] + local_mask_idx
|
|
obj_id = self._object_id(view=view, local_mask_idx=local_mask_idx)
|
|
node = self._create_object_node(
|
|
obj_id=obj_id,
|
|
view=view,
|
|
room=graph.rooms[view.room_id],
|
|
meta=meta,
|
|
selected_idx=selected_idx,
|
|
hash_bits=output.hash_bits,
|
|
crop_index=crop_index,
|
|
)
|
|
graph.objects[obj_id] = node
|
|
artifacts.object_images[obj_id] = output.cropped_images[crop_index]
|
|
|
|
def _object_id(self, *, view: Any, local_mask_idx: int) -> str:
|
|
return f"{view.room_id}_v{view.view_idx:03d}_m{local_mask_idx:02d}"
|
|
|
|
def _source_view_id(self, view: Any) -> str:
|
|
return f"{view.room_id}_v{view.view_idx:03d}"
|
|
|
|
def _create_object_node(
|
|
self,
|
|
*,
|
|
obj_id: str,
|
|
view: Any,
|
|
room: RoomNode,
|
|
meta: dict[str, Any],
|
|
selected_idx: int,
|
|
hash_bits: Any,
|
|
crop_index: int,
|
|
) -> ObjectNode:
|
|
hash_bytes = bits_tensor_to_hash_bytes(hash_bits[crop_index])
|
|
label, confidence, bbox_xyxy = self._metadata_for_detection(
|
|
meta=meta,
|
|
selected_idx=selected_idx,
|
|
)
|
|
return ObjectNode(
|
|
obj_id=obj_id,
|
|
room_id=view.room_id,
|
|
position=room.center.copy(),
|
|
visual_hash=hash_bytes,
|
|
semantic_hash=hash_bytes,
|
|
hit_count=1,
|
|
last_seen_frame=int(view.view_idx),
|
|
label=label,
|
|
confidence=confidence,
|
|
bbox_xyxy=bbox_xyxy,
|
|
source_view_id=self._source_view_id(view),
|
|
position_confidence=0.0,
|
|
)
|
|
|
|
def _metadata_for_detection(
|
|
self,
|
|
*,
|
|
meta: dict[str, Any],
|
|
selected_idx: int,
|
|
) -> tuple[str | None, float | None, tuple[float, float, float, float] | None]:
|
|
label = self._metadata_item(meta=meta, key="labels", selected_idx=selected_idx)
|
|
score = self._metadata_item(meta=meta, key="scores", selected_idx=selected_idx)
|
|
box = self._metadata_item(meta=meta, key="boxes_xyxy", selected_idx=selected_idx)
|
|
|
|
confidence = None if score is None else float(score)
|
|
bbox_xyxy = None if box is None else self._normalize_bbox(box)
|
|
return label, confidence, bbox_xyxy
|
|
|
|
def _metadata_item(self, *, meta: dict[str, Any], key: str, selected_idx: int) -> Any:
|
|
values = meta.get(key)
|
|
if values is None:
|
|
return None
|
|
if not (0 <= selected_idx < len(values)):
|
|
raise ValueError(
|
|
f"selected_idx {selected_idx} out of range for "
|
|
f"metadata key '{key}' with length {len(values)}"
|
|
)
|
|
return values[selected_idx]
|
|
|
|
def _normalize_bbox(self, box: Any) -> tuple[float, float, float, float]:
|
|
if len(box) != 4:
|
|
raise ValueError(f"bbox entry has length {len(box)}, expected 4")
|
|
return tuple(float(x) for x in box)
|