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