Mini-Nav/mini-nav/scenegraph/builder.py

"""Scene graph builder: converts room views + pipeline debug output into SimpleSceneGraph."""

from __future__ import annotations

from dataclasses import dataclass, field
from math import radians, tan
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.
            'room_center' uses the room's center point.
            'bbox_depth_center' uses depth at the bbox center to compute
            a 3D position via pinhole projection.
        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.
        camera_hfov_degrees: Horizontal field of view of the camera in degrees.
            Only used when position_strategy='bbox_depth_center'.
    """

    position_strategy: Literal["room_center", "bbox_depth_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
    camera_hfov_degrees: float = 90.0

    def __post_init__(self) -> None:
        valid_strategies = {"room_center", "bbox_depth_center"}
        if self.position_strategy not in valid_strategies:
            raise ValueError(
                f"position_strategy must be one of {valid_strategies}, "
                f"got {self.position_strategy!r}"
            )
        if not 0 < self.camera_hfov_degrees < 180:
            raise ValueError(
                f"camera_hfov_degrees must be in (0, 180), "
                f"got {self.camera_hfov_degrees}"
            )
        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/M1")


@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,
        )
        position, position_confidence = self._position_for_detection(
            view=view,
            room=room,
            bbox_xyxy=bbox_xyxy,
        )
        return ObjectNode(
            obj_id=obj_id,
            room_id=view.room_id,
            position=position,
            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=position_confidence,
        )

    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] | None:
        try:
            if len(box) != 4:
                return None
            result = tuple(float(x) for x in box)
            if not all(np.isfinite(result)):
                return None
            return result
        except (TypeError, ValueError):
            return None

    def _position_for_detection(
        self,
        *,
        view: Any,
        room: RoomNode,
        bbox_xyxy: tuple[float, float, float, float] | None,
    ) -> tuple[np.ndarray, float]:
        """Compute object position and confidence based on the configured strategy.

        Returns:
            A tuple of (position: np.ndarray (3,), confidence: float).
        """
        if self._config.position_strategy == "room_center":
            return room.center.copy(), 0.0

        # bbox_depth_center strategy
        if bbox_xyxy is not None:
            position = self._compute_bbox_depth_center_position(view, bbox_xyxy)
            if position is not None:
                return position, 1.0

        # Fallback: room center with 0.0 confidence
        return room.center.copy(), 0.0

    def _compute_bbox_depth_center_position(
        self,
        view: Any,
        bbox_xyxy: tuple[float, float, float, float],
    ) -> np.ndarray | None:
        """Compute 3D world position from depth at the bbox centre.

        Uses the configured camera_hfov_degrees for pinhole projection.
        Prefers camera pose (camera_position/camera_rotation) when available
        on the view; falls back to room center + 0.0 confidence when camera
        pose is absent (agent pose alone is insufficient because the Habitat
        depth sensor is offset from the agent body by sensor_height).

        Returns:
            A (3,) float32 position array, or None if computation fails.
        """
        depth = getattr(view, "depth", None)

        if depth is None:
            return None

        # Prefer camera pose when available.  Agent pose alone is insufficient
        # because the Habitat depth sensor is offset [0.0, sensor_height, 0.0]
        # from the agent body, and we have no way to recover the offset.
        camera_position = getattr(view, "camera_position", None)
        camera_rotation = getattr(view, "camera_rotation", None)

        if camera_position is not None and camera_rotation is not None:
            origin_position = camera_position
            origin_rotation = camera_rotation
        else:
            # No camera pose – cannot confidently place in world space.
            return None

        # Coerce depth to ndarray; require 2D
        try:
            depth = np.asarray(depth)
        except (ValueError, TypeError):
            return None
        if depth.ndim != 2:
            return None

        height, width = depth.shape

        # Validate bbox values are finite before rounding
        try:
            x1, y1, x2, y2 = map(float, bbox_xyxy)
        except (ValueError, TypeError):
            return None
        if not all(np.isfinite(v) for v in (x1, y1, x2, y2)):
            return None

        u = int(round((x1 + x2) / 2.0))
        v = int(round((y1 + y2) / 2.0))

        # Out of bounds check
        if not (0 <= u < width and 0 <= v < height):
            return None

        depth_value = depth[v, u]

        # Depth must be finite and positive
        if not np.isfinite(depth_value) or depth_value <= 0:
            return None

        # Convert rotation to matrix
        rotation_matrix = self._rotation_to_matrix(origin_rotation)
        if rotation_matrix is None:
            return None

        # Coerce origin_position safely
        try:
            origin_pos = np.asarray(origin_position, dtype=np.float32).reshape(3)
        except (ValueError, TypeError, RuntimeError):
            return None
        if not np.all(np.isfinite(origin_pos)):
            return None

        # Pinhole projection
        camera_point = self._pixel_depth_to_camera_point(
            u=u, v=v, depth=depth_value, width=width, height=height,
        )

        # World position
        world_position = origin_pos + rotation_matrix @ camera_point

        if not np.all(np.isfinite(world_position)):
            return None

        return np.asarray(world_position, dtype=np.float32)

    def _pixel_depth_to_camera_point(
        self,
        u: int,
        v: int,
        depth: float,
        width: int,
        height: int,
    ) -> np.ndarray:
        """Convert a pixel with depth to a camera-space 3D point via pinhole model.

        M1 assumes Habitat camera forward is local -Z; this convention should be
        validated with a real Habitat sanity check.

        Args:
            u: Pixel column (x) coordinate.
            v: Pixel row (y) coordinate.
            depth: Depth value at the pixel (positive, finite).
            width: Image width in pixels.
            height: Image height in pixels.

        Returns:
            A (3,) float32 array [x, y, z] in camera space.
        """
        hfov = radians(self._config.camera_hfov_degrees)
        fx = width / (2.0 * tan(hfov / 2.0))
        fy = fx
        cx = (width - 1) / 2.0
        cy = (height - 1) / 2.0

        x = (u - cx) / fx * depth
        y = -(v - cy) / fy * depth
        z = -depth

        return np.array([x, y, z], dtype=np.float32)

    def _rotation_to_matrix(
        self,
        rotation: Any,
    ) -> np.ndarray | None:
        """Convert a rotation object to a (3, 3) float32 rotation matrix.

        Supports:
        - Objects with a ``.rotation_matrix`` attribute (e.g. habitat_sim style).
        - ndarray or array-like objects directly convertible to (3, 3).
        - ``numpy-quaternion`` quaternion objects (optional dependency; uses
          ``importlib`` to load ``quaternion.as_rotation_matrix``).
        - Objects with a ``.transform_vector()`` method (quaternion-like);
          constructs the matrix columns by transforming local basis vectors.

        Returns:
            A (3, 3) float32 matrix, or None if conversion fails.
        """
        try:
            # Path a: object with .rotation_matrix attribute
            if hasattr(rotation, "rotation_matrix"):
                mat = self._validate_matrix(rotation.rotation_matrix)
                if mat is not None:
                    return mat
                return None

            # Path b: direct ndarray or array-like (3, 3)
            mat = self._validate_matrix(rotation)
            if mat is not None:
                return mat

            # Path c: numpy-quaternion quaternion object (optional dependency)
            try:
                import importlib as _il
                qmod = _il.import_module("quaternion")
                if hasattr(qmod, "as_rotation_matrix"):
                    mat = self._validate_matrix(qmod.as_rotation_matrix(rotation))
                    if mat is not None:
                        return mat
            except (ImportError, AttributeError, TypeError, ValueError):
                pass

            # Path d: quaternion-like with .transform_vector method
            if hasattr(rotation, "transform_vector"):
                eye = np.eye(3, dtype=np.float32)
                cols: list[np.ndarray] = []
                for i in range(3):
                    col_vec = rotation.transform_vector(eye[:, i])
                    col = np.asarray(col_vec, dtype=np.float32).ravel()
                    if col.shape != (3,) or not np.all(np.isfinite(col)):
                        return None
                    cols.append(col)
                return np.column_stack(cols)

            return None
        except (ValueError, TypeError, RuntimeError):
            return None

    @staticmethod
    def _validate_matrix(mat: Any) -> np.ndarray | None:
        """Validate and convert a candidate to a (3, 3) float32 matrix.

        Returns the matrix if it is (3, 3) with all finite values, else None.
        """
        try:
            mat = np.asarray(mat, dtype=np.float32)
            if mat.shape == (3, 3) and np.all(np.isfinite(mat)):
                return mat
        except (ValueError, TypeError, RuntimeError):
            pass
        return None