diff --git a/mini-nav/scenegraph/__init__.py b/mini-nav/scenegraph/__init__.py index b4ec631..987f667 100644 --- a/mini-nav/scenegraph/__init__.py +++ b/mini-nav/scenegraph/__init__.py @@ -5,6 +5,7 @@ This module exports the main scenegraph objects for easy import: from scenegraph import SimpleSceneGraph, RoomNode, ObjectNode """ +from .builder import SceneGraphBuildArtifacts, SceneGraphBuildConfig, SceneGraphBuilder from .hash_codec import ( bits_tensor_to_cam_row, bits_tensor_to_hash_bytes, @@ -28,6 +29,9 @@ __all__ = [ "ImageSceneGraphQueryResult", "ObjectNode", "RoomNode", + "SceneGraphBuildArtifacts", + "SceneGraphBuildConfig", + "SceneGraphBuilder", "SceneGraphMatch", "SimpleSceneGraph", "SoftwareCamIndex", diff --git a/mini-nav/scenegraph/builder.py b/mini-nav/scenegraph/builder.py new file mode 100644 index 0000000..3af53a6 --- /dev/null +++ b/mini-nav/scenegraph/builder.py @@ -0,0 +1,293 @@ +"""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) diff --git a/mini-nav/scenegraph/objectnode.py b/mini-nav/scenegraph/objectnode.py index b1ba4c7..863e231 100644 --- a/mini-nav/scenegraph/objectnode.py +++ b/mini-nav/scenegraph/objectnode.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from dataclasses import dataclass import numpy as np @@ -20,10 +22,20 @@ class ObjectNode: hit_count: int = 1 # 被观测到的次数。太低的可以直接过滤掉 last_seen_frame: int = 0 # 最后一次看到的帧号或时间戳 + # Optional detection metadata (from M0 integration) + label: str | None = None + confidence: float | None = None + bbox_xyxy: tuple[float, float, float, float] | None = None + source_view_id: str | None = None + position_confidence: float | None = None + def __post_init__(self): + self.position = np.asarray(self.position, dtype=np.float32) if len(self.visual_hash) != 64: raise ValueError("visual_hash must be exactly 64 bytes (512 bits)") if len(self.semantic_hash) != 64: raise ValueError("semantic_hash must be exactly 64 bytes (512 bits)") if self.position.shape != (3,): raise ValueError("position must have shape (3,)") + if self.bbox_xyxy is not None and len(self.bbox_xyxy) != 4: + raise ValueError("bbox_xyxy must have exactly 4 elements") diff --git a/mini-nav/simulator/__init__.py b/mini-nav/simulator/__init__.py index 170f74f..9e4c950 100644 --- a/mini-nav/simulator/__init__.py +++ b/mini-nav/simulator/__init__.py @@ -5,10 +5,17 @@ from .habitat import ( ) from .image_save import save_object_image, save_room_view from .topdown import TopDownRenderStyle, TopDownSceneElements, render_topdown_scene_map -from .views import RoomViewsByRoom, collect_room_views_by_room, collect_scene_images +from .views import ( + RoomView, + RoomViewsByRoom, + collect_room_views_by_room, + collect_scene_images, + flatten_room_views, +) __all__ = [ "HabitatSimulatorConfig", + "RoomView", "RoomViewsByRoom", "TopDownRenderStyle", "TopDownSceneElements", @@ -16,6 +23,7 @@ __all__ = [ "collect_room_views_by_room", "collect_scene_images", "create_habitat_simulator", + "flatten_room_views", "render_topdown_scene_map", "save_object_image", "save_room_view", diff --git a/mini-nav/simulator/views.py b/mini-nav/simulator/views.py index d61b13b..8fefbd5 100644 --- a/mini-nav/simulator/views.py +++ b/mini-nav/simulator/views.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import dataclass, field from importlib import import_module from pathlib import Path from typing import Any, Callable, Iterable, Sequence @@ -11,6 +12,24 @@ RoomViewsByRoom = dict[str, list[Any]] ProgressTrack = Callable[[Iterable[Any], str], Iterable[Any]] +@dataclass(frozen=True) +class RoomView: + room_id: str + view_idx: int + rgb: Any = field(compare=False) + depth: Any | None = field(default=None, compare=False) + agent_position: np.ndarray | None = field(default=None, compare=False) + agent_rotation: Any | None = field(default=None, compare=False) + + +def flatten_room_views(room_views_by_room: RoomViewsByRoom) -> list[RoomView]: + result: list[RoomView] = [] + for room_id, views in room_views_by_room.items(): + for idx, rgb in enumerate(views): + result.append(RoomView(room_id=room_id, view_idx=idx, rgb=rgb)) + return result + + def collect_room_views_by_room( agent: Any, sim: Any, diff --git a/notebooks/verification.py b/notebooks/verification.py index 98ec317..b53ed14 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -36,10 +36,9 @@ def project_imports(): from compressors import HashPipeline from configs import cfg_manager from scenegraph import ( - ObjectNode, RoomNode, - SimpleSceneGraph, - bits_tensor_to_hash_bytes, + SceneGraphBuildConfig, + SceneGraphBuilder, query_image_against_scene_graph, ) from simulator import ( @@ -47,6 +46,7 @@ def project_imports(): TopDownSceneElements, collect_room_views_by_room, create_habitat_simulator, + flatten_room_views, render_topdown_scene_map, save_object_image, save_room_view, @@ -56,14 +56,14 @@ def project_imports(): return ( HashPipeline, HabitatSimulatorConfig, - ObjectNode, RoomNode, - SimpleSceneGraph, + SceneGraphBuildConfig, + SceneGraphBuilder, TopDownSceneElements, - bits_tensor_to_hash_bytes, cfg_manager, collect_room_views_by_room, create_habitat_simulator, + flatten_room_views, numpy_to_pil, query_image_against_scene_graph, render_topdown_scene_map, @@ -163,6 +163,7 @@ def pipeline_init(HashPipeline): def collect_views( agent, collect_room_views_by_room, + flatten_room_views, numpy_to_pil, room_nodes, sim, @@ -175,162 +176,108 @@ def collect_views( views_per_room=views_per_room, ) - # Flatten room views into (room_id, view_idx, PIL.Image) tuples. + room_views = flatten_room_views(all_room_views) + + # Build dataset of (room_id, view_idx, PIL.Image) tuples. room_view_dataset = [ - (_room_id, _view_idx, numpy_to_pil(_rgb)) - for _room_id, _views in all_room_views.items() - for _view_idx, _rgb in enumerate(_views) + (_view.room_id, _view.view_idx, numpy_to_pil(_view.rgb)) + for _view in room_views ] print(f"Collected {len(room_view_dataset)} room views") - return all_room_views, room_view_dataset + return all_room_views, room_view_dataset, room_views @app.cell def build_scene_graph( - ObjectNode, - SimpleSceneGraph, - bits_tensor_to_hash_bytes, + SceneGraphBuildConfig, + SceneGraphBuilder, cfg_manager, mo, - np, + numpy_to_pil, pipeline, room_nodes, - room_view_dataset, + room_views, save_object_image, save_room_view, text_labels, - torch, ): - scene_graph = SimpleSceneGraph( - rooms={_room.room_id: _room for _room in room_nodes}, - objects={}, - ) - - # Storage for cropped object images (for visualization). - object_images = {} - output_dir = cfg_manager.get().output.directory / "verification" output_dir.mkdir(parents=True, exist_ok=True) - _images = [item[2] for item in room_view_dataset] - _metadata = [(item[0], item[1]) for item in room_view_dataset] + builder = SceneGraphBuilder( + pipeline=pipeline, + config=SceneGraphBuildConfig(inference_batch_size=4), + ) - inference_batch_size = 4 - image_batches = [ - _images[index : index + inference_batch_size] - for index in range(0, len(_images), inference_batch_size) + pil_room_views = [ + type(_view)( + room_id=_view.room_id, + view_idx=_view.view_idx, + rgb=numpy_to_pil(_view.rgb), + depth=_view.depth, + agent_position=_view.agent_position, + agent_rotation=_view.agent_rotation, + ) + for _view in room_views ] - _cropped_images = [] - debug_meta = [] - hash_batches = [] - for _batch_images in mo.status.progress_bar( - image_batches, - title="Running pipeline inference on room views", - subtitle=f"Batch size {inference_batch_size} with ETA", - completion_title="Pipeline inference finished", - completion_subtitle=( - f"Processed {len(_images)} room views in {len(image_batches)} batches" - ), + with mo.status.spinner(title="Building scene graph from room views"): + scene_graph, build_artifacts = builder.build_from_room_views( + room_nodes=room_nodes, + room_views=pil_room_views, + text_labels=text_labels, + ) + + object_images = build_artifacts.object_images + debug_meta = build_artifacts.debug_meta + + # Save original room views. + for _room_view in mo.status.progress_bar( + pil_room_views, + title="Saving room-view snapshots", + subtitle="Writing original room images to disk", + completion_title="Room-view snapshots saved", + completion_subtitle=f"Saved {len(pil_room_views)} room views", show_eta=True, show_rate=True, remove_on_exit=False, ): - _batch_output = pipeline.process_batch( - _batch_images, - text_labels, - batch_size=inference_batch_size, - return_debug_details=True, - ) - _cropped_images.extend(_batch_output.cropped_images) - debug_meta.extend(_batch_output.debug_meta) - if _batch_output.hash_bits.numel() > 0: - hash_batches.append(_batch_output.hash_bits) + save_room_view(output_dir, _room_view.room_id, _room_view.view_idx, _room_view.rgb) - if hash_batches: - hash_tensor = torch.cat(hash_batches, dim=0) - else: - hash_tensor = torch.empty( - (0, pipeline.hash_bits), dtype=torch.int32, device=pipeline.device + # Save object crops. + for _obj_id, _cropped in mo.status.progress_bar( + object_images.items(), + title="Saving object crops", + subtitle="Writing cropped object images to disk", + completion_title="Object crops saved", + completion_subtitle=f"Saved {len(object_images)} object crops", + show_eta=True, + show_rate=True, + remove_on_exit=False, + ): + _node = scene_graph.objects[_obj_id] + save_object_image( + output_dir, + _node.room_id, + _obj_id, + _node.last_seen_frame, + 0, # mask idx 0 ok for M0 + _cropped, ) from collections import Counter - _reasons = Counter(m["fallback_reason"] or "ok" for m in debug_meta) + _meta_fallbacks = [_meta.get("fallback_reason") for _meta in debug_meta] + fallback_count = sum(1 for f in _meta_fallbacks if f is not None) + _reasons = Counter(f or "ok" for f in _meta_fallbacks) print(f"Fallback breakdown: {_reasons}") - # Save original room views. - for _room_id, _view_idx, _image in mo.status.progress_bar( - room_view_dataset, - title="Saving room-view snapshots", - subtitle="Writing original room images to disk", - completion_title="Room-view snapshots saved", - completion_subtitle=f"Saved {len(room_view_dataset)} room views", - show_eta=True, - show_rate=True, - remove_on_exit=False, - ): - save_room_view(output_dir, _room_id, _view_idx, _image) - - # Prefix sum: map flat crop index to (input_image_idx, mask_idx). - _num_selected = [_m["num_selected"] for _m in debug_meta] - assert sum(_num_selected) == len(_cropped_images), ( - f"Sum of num_selected ({sum(_num_selected)}) != cropped_images count ({len(_cropped_images)})" - ) - _prefix_sums = [0] - for _n in _num_selected: - _prefix_sums.append(_prefix_sums[-1] + _n) - - _obj_counter = 0 - _total_crops = len(_cropped_images) - object_tasks = [] - for _img_idx, _n_crops in enumerate(_num_selected): - _room_id, _view_idx = _metadata[_img_idx] - for _mask_idx in range(_n_crops): - object_tasks.append((_img_idx, _room_id, _view_idx, _mask_idx, _n_crops)) - - for _img_idx, _room_id, _view_idx, _mask_idx, _n_crops in mo.status.progress_bar( - object_tasks, - title="Building scene graph objects", - subtitle="Preparing cropped objects and hashes with ETA", - completion_title="Scene graph build complete", - completion_subtitle=f"Created {_total_crops} cropped object entries", - show_eta=True, - show_rate=True, - remove_on_exit=False, - ): - _crop_flat_idx = _prefix_sums[_img_idx] + _mask_idx - _cropped = _cropped_images[_crop_flat_idx] - _hash_bits = hash_tensor[_crop_flat_idx] - - _obj_id = f"{_room_id}_v{_view_idx:03d}_m{_mask_idx:02d}" - - _hash_bytes = bits_tensor_to_hash_bytes(_hash_bits) - - object_images[_obj_id] = _cropped - save_object_image(output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped) - - scene_graph.objects[_obj_id] = ObjectNode( - obj_id=_obj_id, - room_id=_room_id, - position=np.array([0.0, 0.0, 0.0], dtype=np.float32), - visual_hash=_hash_bytes, - semantic_hash=_hash_bytes, - hit_count=1, - last_seen_frame=_view_idx, - ) - _obj_counter += 1 - - _fallback_count = sum( - 1 for _meta in debug_meta if _meta["fallback_reason"] is not None - ) - - print(f"Created {_obj_counter} objects") + print(f"Created {len(scene_graph.objects)} objects") print(f"Saved cropped images to: {output_dir}") - print(f"Fallback frames: {_fallback_count}/{len(debug_meta)}") + print(f"Fallback frames: {fallback_count}/{len(debug_meta)}") - return hash_tensor, object_images, output_dir, scene_graph + return build_artifacts, object_images, output_dir, scene_graph @app.cell @@ -350,6 +297,12 @@ def build_tables(pl, scene_graph): "obj_id": obj.obj_id, "room_id": obj.room_id, "visual_hash": obj.visual_hash.hex()[:16] + "...", + "label": obj.label, + "confidence": obj.confidence, + "source_view_id": obj.source_view_id, + "position_x": float(obj.position[0]) if obj.position is not None else None, + "position_y": float(obj.position[1]) if obj.position is not None else None, + "position_z": float(obj.position[2]) if obj.position is not None else None, } for obj in scene_graph.objects.values() ] diff --git a/tests/test_notebook_verification_static.py b/tests/test_notebook_verification_static.py index ddfa7f4..2aea43d 100644 --- a/tests/test_notebook_verification_static.py +++ b/tests/test_notebook_verification_static.py @@ -16,5 +16,18 @@ def test_verification_notebook_uses_scenegraph_query_api(): def test_verification_notebook_uses_hash_codec_for_object_hashes(): source = NOTEBOOK.read_text() - assert "bits_tensor_to_hash_bytes" in source + # Notebook should no longer do hash conversion directly; + # the builder handles it internally. + assert "bits_tensor_to_hash_bytes" not in source assert "np.packbits" not in source + assert "scene_graph.objects[_obj_id] = ObjectNode" not in source + + +def test_verification_notebook_uses_scene_graph_builder(): + source = NOTEBOOK.read_text() + + assert "SceneGraphBuilder" in source + assert "SceneGraphBuildConfig" in source + assert "flatten_room_views" in source + assert "scene_graph, build_artifacts = builder.build_from_room_views" in source + assert "scene_graph.objects[_obj_id] = ObjectNode" not in source diff --git a/tests/test_scenegraph_builder.py b/tests/test_scenegraph_builder.py new file mode 100644 index 0000000..86a75d7 --- /dev/null +++ b/tests/test_scenegraph_builder.py @@ -0,0 +1,479 @@ +"""Tests for scenegraph builder / ObjectNode.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +MINI_NAV_DIR = Path(__file__).resolve().parents[1] / "mini-nav" +sys.path.insert(0, str(MINI_NAV_DIR)) + +from scenegraph.hash_codec import bits_tensor_to_hash_bytes # noqa: E402 +from scenegraph.objectnode import ObjectNode # noqa: E402 + + +def _hash_bytes() -> bytes: + """Create a 512-bit hash with bit 0 set to 1.""" + bits = torch.zeros(512, dtype=torch.uint8) + bits[0] = 1 + return bits_tensor_to_hash_bytes(bits) + + +def test_object_node_accepts_optional_detection_metadata(): + """ObjectNode should accept and preserve optional detection metadata.""" + hb = _hash_bytes() + node = ObjectNode( + obj_id="room_a_v000_m00", + room_id="room_a", + position=np.array([1, 2, 3], dtype=np.float32), + visual_hash=hb, + semantic_hash=hb, + hit_count=1, + last_seen_frame=0, + label="a chair", + confidence=0.87, + bbox_xyxy=(10, 20, 30, 40), + source_view_id="room_a_v000", + position_confidence=0.0, + ) + + assert node.obj_id == "room_a_v000_m00" + assert node.room_id == "room_a" + np.testing.assert_array_equal(node.position, np.array([1, 2, 3], dtype=np.float32)) + assert node.visual_hash == hb + assert node.semantic_hash == hb + assert node.hit_count == 1 + assert node.last_seen_frame == 0 + assert node.label == "a chair" + assert node.confidence == 0.87 + assert node.bbox_xyxy == (10, 20, 30, 40) + assert node.source_view_id == "room_a_v000" + assert node.position_confidence == 0.0 + + +def test_room_view_equality_with_numpy_rgb_does_not_compare_payload(): + """RoomView equality should not attempt element-wise numpy comparison.""" + from simulator import RoomView # noqa: E402 + + rgb_a = np.zeros((2, 2, 3), dtype=np.uint8) + rgb_b = np.ones((2, 2, 3), dtype=np.uint8) + + # Two views with different numpy payloads — still equal because + # rgb is excluded from equality (only room_id, view_idx matter). + v1 = RoomView(room_id="room_a", view_idx=0, rgb=rgb_a) + v2 = RoomView(room_id="room_a", view_idx=0, rgb=rgb_b) + assert v1 == v2 + + # Different room_id → not equal + v3 = RoomView(room_id="room_b", view_idx=0, rgb=rgb_a) + assert v1 != v3 + + # Different view_idx → not equal + v4 = RoomView(room_id="room_a", view_idx=1, rgb=rgb_a) + assert v1 != v4 + + +def test_flatten_room_views_with_numpy_arrays_safe_equality(): + """flatten_room_views with numpy arrays should not raise ambiguous truth value errors.""" + from simulator import RoomView, flatten_room_views # noqa: E402 + + rgb_a = np.zeros((2, 2, 3), dtype=np.uint8) + rgb_b = np.ones((2, 2, 3), dtype=np.uint8) + + views = flatten_room_views({"room_a": [rgb_a, rgb_b]}) + + # Compare with expected list — uses identity-field comparison only + expected = [ + RoomView(room_id="room_a", view_idx=0, rgb=rgb_a), + RoomView(room_id="room_a", view_idx=1, rgb=rgb_b), + ] + assert views == expected + + # Even with wildly different numpy arrays, comparison is safe + swapped = [ + RoomView(room_id="room_a", view_idx=0, rgb=rgb_b), + RoomView(room_id="room_a", view_idx=1, rgb=rgb_a), + ] + assert views == swapped # rgb payload ignored + + # Mismatched identity fields still differ + different_room = [ + RoomView(room_id="room_b", view_idx=0, rgb=rgb_a), + RoomView(room_id="room_a", view_idx=1, rgb=rgb_b), + ] + assert views != different_room + + +def test_flatten_room_views_preserves_room_and_view_identity(): + from PIL import Image # noqa: E402 + from simulator import RoomView, flatten_room_views # noqa: E402 + + first = Image.new("RGB", (4, 4), "red") + second = Image.new("RGB", (4, 4), "blue") + views = flatten_room_views({"room_a": [first, second]}) + + # Ensure the returned objects are the same instances + assert views[0].rgb is first + assert views[1].rgb is second + + # Equality compares only room_id and view_idx (payload excluded). + # Different PIL images with same room_id/view_idx compare equal. + assert views == [ + RoomView(room_id="room_a", view_idx=0, rgb=first), + RoomView(room_id="room_a", view_idx=1, rgb=second), + ] + assert views == [ + RoomView(room_id="room_a", view_idx=0, rgb=second), # rgb ignored + RoomView(room_id="room_a", view_idx=1, rgb=first), # rgb ignored + ] + + # Different room_id or view_idx makes them unequal + assert views != [ + RoomView(room_id="room_b", view_idx=0, rgb=first), + RoomView(room_id="room_a", view_idx=1, rgb=second), + ] + + +# --------------------------------------------------------------------------- +# SceneGraphBuilder M0 tests +# --------------------------------------------------------------------------- + + +class FakePipeline: + """Fake pipeline that records calls and returns a pre-set output.""" + + def __init__(self, output, calls, hash_bits=512): + self.output = output + self.calls = calls + self._hash_bits = hash_bits + + def process_batch(self, images, text_labels, batch_size=32, return_debug_details=False): + self.calls.append({ + "images": images, + "text_labels": text_labels, + "batch_size": batch_size, + "return_debug_details": return_debug_details, + }) + return self.output + + +def test_scene_graph_builder_creates_objects_from_pipeline_debug_meta(): + """Build a graph from one room view with one detection and verify all fields.""" + from scenegraph import RoomNode, SceneGraphBuilder, SceneGraphBuildConfig + from simulator import RoomView + + room_node = RoomNode(room_id="room_a", center=[10, 0, 20], bbox_extent=[5, 3, 5]) + room_nodes = [room_node] + + rgb = np.zeros((32, 32, 3), dtype=np.uint8) + room_view = RoomView(room_id="room_a", view_idx=0, rgb=rgb) + room_views = [room_view] + + text_labels = ["a chair"] + + hash_bits = torch.zeros(1, 512, dtype=torch.uint8) + hash_bits[0, 3] = 1 + + crop = "crop_data_0" + + debug_meta_entry = { + "selected_indices": [0], + "boxes_xyxy": [[10.0, 20.0, 30.0, 40.0]], + "scores": [0.91], + "labels": ["a chair"], + "masks": [np.zeros((32, 32), dtype=np.uint8)], + "num_selected": 1, + } + + output = SimpleNamespace( + cropped_images=[crop], + hash_bits=hash_bits, + debug_meta=[debug_meta_entry], + ) + + pipeline = FakePipeline(output=output, calls=[], hash_bits=512) + config = SceneGraphBuildConfig(inference_batch_size=7) + builder = SceneGraphBuilder(pipeline=pipeline, config=config) + + graph, artifacts = builder.build_from_room_views( + room_nodes=room_nodes, + room_views=room_views, + text_labels=text_labels, + ) + + # Pipeline was called correctly + assert len(pipeline.calls) == 1 + call = pipeline.calls[0] + assert call["return_debug_details"] is True + assert call["batch_size"] == 7 + assert len(call["images"]) == 1 + assert call["images"][0] is rgb + assert call["text_labels"] == text_labels + + # Graph rooms + assert list(graph.rooms.keys()) == ["room_a"] + assert graph.rooms["room_a"] is room_node + + # Graph objects + assert list(graph.objects.keys()) == ["room_a_v000_m00"] + + obj = graph.objects["room_a_v000_m00"] + assert obj.room_id == "room_a" + np.testing.assert_array_equal(obj.position, np.array([10, 0, 20], dtype=np.float32)) + + expected_hash = bits_tensor_to_hash_bytes(hash_bits[0]) + assert obj.visual_hash == expected_hash + assert obj.semantic_hash == expected_hash + + assert obj.label == "a chair" + assert obj.confidence == 0.91 + assert obj.bbox_xyxy == (10.0, 20.0, 30.0, 40.0) + assert obj.source_view_id == "room_a_v000" + assert obj.position_confidence == 0.0 + + # Artifacts + assert artifacts.object_images["room_a_v000_m00"] is crop + assert artifacts.debug_meta == [debug_meta_entry] + + +def test_scene_graph_builder_handles_empty_pipeline_output(): + """Pipeline returning no detections should produce rooms-only graph.""" + from scenegraph import RoomNode, SceneGraphBuilder + from simulator import RoomView + + room_node = RoomNode(room_id="room_a", center=[0, 0, 0], bbox_extent=[5, 3, 5]) + room_nodes = [room_node] + + rgb = np.zeros((32, 32, 3), dtype=np.uint8) + room_view = RoomView(room_id="room_a", view_idx=0, rgb=rgb) + room_views = [room_view] + + text_labels: list[str] = [] + + hash_bits = torch.zeros(0, 512, dtype=torch.uint8) + + debug_meta_entry = { + "selected_indices": [], + "boxes_xyxy": [], + "scores": [], + "labels": [], + "masks": [], + "num_selected": 0, + } + + output = SimpleNamespace( + cropped_images=[], + hash_bits=hash_bits, + debug_meta=[debug_meta_entry], + ) + + pipeline = FakePipeline(output=output, calls=[]) + builder = SceneGraphBuilder(pipeline=pipeline) + + graph, artifacts = builder.build_from_room_views( + room_nodes=room_nodes, + room_views=room_views, + text_labels=text_labels, + ) + + # Only rooms, no objects + assert list(graph.rooms.keys()) == ["room_a"] + assert graph.objects == {} + assert artifacts.object_images == {} + + +def test_config_rejects_enable_fusion(): + """SceneGraphBuildConfig should reject enable_fusion=True.""" + from scenegraph import SceneGraphBuildConfig + + with pytest.raises(ValueError, match="fusion"): + SceneGraphBuildConfig(enable_fusion=True) + + +def test_config_rejects_non_room_center_strategy(): + """SceneGraphBuildConfig should reject non-room_center position_strategy.""" + from scenegraph import SceneGraphBuildConfig + + with pytest.raises(ValueError, match="room_center"): + SceneGraphBuildConfig(position_strategy="global") + + +def test_config_rejects_zero_batch_size(): + """SceneGraphBuildConfig should reject inference_batch_size <= 0.""" + from scenegraph import SceneGraphBuildConfig + + with pytest.raises(ValueError): + SceneGraphBuildConfig(inference_batch_size=0) + + +def test_scene_graph_builder_public_method_stays_small(): + """build_from_room_views should stay an orchestration method, not own all details.""" + import inspect + + from scenegraph import SceneGraphBuilder + + source = inspect.getsource(SceneGraphBuilder.build_from_room_views) + code_lines = [ + line for line in source.splitlines() + if line.strip() and not line.strip().startswith("#") + ] + + assert len(code_lines) <= 45 + for helper_name in [ + "_prepare_graph", + "_run_pipeline", + "_build_artifacts", + "_validate_pipeline_output", + "_add_objects_to_graph", + ]: + assert helper_name in source + + +# --------------------------------------------------------------------------- +# Validation error tests +# --------------------------------------------------------------------------- + + +def test_builder_raises_on_selected_indices_num_selected_mismatch(): + """selected_indices length != num_selected should raise ValueError.""" + from scenegraph import RoomNode, SceneGraphBuilder + from simulator import RoomView + + room_node = RoomNode(room_id="room_a", center=[0, 0, 0], bbox_extent=[5, 3, 5]) + room_nodes = [room_node] + rgb = np.zeros((32, 32, 3), dtype=np.uint8) + room_view = RoomView(room_id="room_a", view_idx=0, rgb=rgb) + room_views = [room_view] + + debug_meta_entry = { + "selected_indices": [0], # len 1 + "boxes_xyxy": [[10.0, 20.0, 30.0, 40.0]], + "scores": [0.91], + "labels": ["a chair"], + "num_selected": 0, # mismatch! + } + + output = SimpleNamespace( + cropped_images=["crop"], + hash_bits=torch.zeros(1, 512, dtype=torch.uint8), + debug_meta=[debug_meta_entry], + ) + + pipeline = FakePipeline(output=output, calls=[]) + builder = SceneGraphBuilder(pipeline=pipeline) + + with pytest.raises(ValueError, match="selected_indices"): + builder.build_from_room_views( + room_nodes=room_nodes, + room_views=room_views, + text_labels=[], + ) + + +def test_builder_raises_on_selected_idx_out_of_labels_range(): + """selected_idx >= len(labels) should raise ValueError.""" + from scenegraph import RoomNode, SceneGraphBuilder + from simulator import RoomView + + room_node = RoomNode(room_id="room_a", center=[0, 0, 0], bbox_extent=[5, 3, 5]) + room_nodes = [room_node] + rgb = np.zeros((32, 32, 3), dtype=np.uint8) + room_view = RoomView(room_id="room_a", view_idx=0, rgb=rgb) + room_views = [room_view] + + debug_meta_entry = { + "selected_indices": [5], # index 5, but labels has only 1 entry + "boxes_xyxy": [[10.0, 20.0, 30.0, 40.0]], + "scores": [0.91], + "labels": ["a chair"], + "num_selected": 1, + } + + output = SimpleNamespace( + cropped_images=["crop"], + hash_bits=torch.zeros(1, 512, dtype=torch.uint8), + debug_meta=[debug_meta_entry], + ) + + pipeline = FakePipeline(output=output, calls=[]) + builder = SceneGraphBuilder(pipeline=pipeline) + + with pytest.raises(ValueError, match="selected_idx.*labels"): + builder.build_from_room_views( + room_nodes=room_nodes, + room_views=room_views, + text_labels=[], + ) + + +def test_builder_raises_on_missing_room_before_pipeline(): + """Unknown room_id in room_view should raise ValueError before pipeline call.""" + from scenegraph import RoomNode, SceneGraphBuilder + from simulator import RoomView + + room_node = RoomNode(room_id="room_a", center=[0, 0, 0], bbox_extent=[5, 3, 5]) + room_nodes = [room_node] + rgb = np.zeros((32, 32, 3), dtype=np.uint8) + room_view = RoomView(room_id="room_b", view_idx=0, rgb=rgb) # not in room_nodes! + room_views = [room_view] + + output = SimpleNamespace( + cropped_images=[], + hash_bits=torch.zeros(0, 512, dtype=torch.uint8), + debug_meta=[], + ) + + pipeline = FakePipeline(output=output, calls=[]) + builder = SceneGraphBuilder(pipeline=pipeline) + + with pytest.raises(ValueError, match="room_b"): + builder.build_from_room_views( + room_nodes=room_nodes, + room_views=room_views, + text_labels=[], + ) + + # Pipeline should NOT have been called — error raised before pipeline + assert len(pipeline.calls) == 0 + + +def test_builder_raises_on_debug_meta_length_mismatch(): + """Mismatched debug_meta length should raise ValueError.""" + from scenegraph import RoomNode, SceneGraphBuilder + from simulator import RoomView + + room_node = RoomNode(room_id="room_a", center=[0, 0, 0], bbox_extent=[5, 3, 5]) + room_nodes = [room_node] + rgb = np.zeros((32, 32, 3), dtype=np.uint8) + room_view = RoomView(room_id="room_a", view_idx=0, rgb=rgb) + room_views = [room_view] + + debug_meta_entry = { + "selected_indices": [], + "boxes_xyxy": [], + "scores": [], + "labels": [], + "num_selected": 0, + } + + # 2 debug_meta entries but only 1 room_view + output = SimpleNamespace( + cropped_images=[], + hash_bits=torch.zeros(0, 512, dtype=torch.uint8), + debug_meta=[debug_meta_entry, debug_meta_entry], + ) + + pipeline = FakePipeline(output=output, calls=[]) + builder = SceneGraphBuilder(pipeline=pipeline) + + with pytest.raises(ValueError, match="debug_meta"): + builder.build_from_room_views( + room_nodes=room_nodes, + room_views=room_views, + text_labels=[], + )