src/fpgm/objects/types.py

"""Shared contracts for the object reconstruction + interaction pipeline.

The pipeline is deliberately staged so every step is separately runnable and
leaves an inspectable artefact on disk:

    1. mask   : SAM 3.1 mask -> tight RGBA crop of the manipulated object
    2. mesh   : crop -> watertight textured mesh (TRELLIS.2, or a fitted proxy)
    3. align  : mask x depth -> world point cloud -> scale + 6-DOF mesh pose
    4. act    : robot FK -> per-frame object pose (free / pushed / grasped)
    5. render : point cloud + robot + object, in the world frame

Every stage consumes and returns the dataclasses below, so a stage can be
re-run from the previous stage's saved artefact without repeating the ones
before it -- which is what makes the thing debuggable.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any

import numpy as np


class MeshSource(Enum):
    """Where an :class:`ObjectMesh` came from.

    Carried through to the render so a proxy is never silently mistaken for a
    real reconstruction.
    """

    TRELLIS = "trellis"
    #: facebook/sam-3d-objects (SAM 3D Objects). Its pipeline decodes a
    #: structured latent to *both* a Gaussian splat and a triangle mesh, and
    #: bakes the mesh's texture from the splat -- see
    #: scripts/sam3d_generate.py and sam3d_objects.pipeline.inference_pipeline
    #: .postprocess_slat_output/to_glb in the vendored repo for the exact path.
    SAM3D = "sam3d"
    #: Oriented bounding box fitted to the observed point cloud. Lets stages
    #: 3-5 run end-to-end while the generator is unavailable. Also the fix
    #: for the drawer/shelf cluster's mesh (Defect 4, see
    #: ``scripts/run_datagen.py``'s ``_build_drawer_slab_mesh``): a drawer
    #: front is a MuJoCo *slide-joint panel*, not a watertight solid, and the
    #: shelf/tray behind it is *background*, not a second object -- fed only
    #: the front panel's own points (separated from the static shelf/tray by
    #: which points actually translate over the episode, gated on the
    #: cluster's own PnP track first passing a mesh-free prismatic check), a
    #: thin oriented box is a faithful "slab" model with no cavity to
    #: swallow anything resting near it.
    PROXY_BOX = "proxy_box"
    #: Convex hull of the observed points -- closer to the real silhouette than
    #: a box, still no invented geometry. **Tried for the drawer/shelf
    #: cluster and rejected twice over** (see
    #: ``scripts/run_datagen.py``'s own module-level comment on Defect 4):
    #: first as a plain hull (solid where the real object is concave, ~2x
    #: silhouette over-coverage, fully occludes an object resting in the
    #: real cavity), then as a grid-triangulated concave depth-surface
    #: reconstruction (fixed the occlusion bug but was still solving the
    #: wrong problem -- see PROXY_BOX's own docstring for the actual fix:
    #: the "drawer" is two different kinds of thing conflated into one S4
    #: cluster, not one oddly-shaped object).
    PROXY_HULL = "proxy_hull"
    #: A directly observed point cloud (``faces`` is ``(0, 3)`` -- no
    #: triangles, ever), not a fitted or invented shape of any kind.
    #: **The drawer's third and current representation** (Defect 4,
    #: ``scripts/run_datagen.py``'s own module-level comment): ``PROXY_BOX``
    #: (a slab, itself replacing ``PROXY_HULL``'s two rejected attempts) was
    #: also measurably wrong -- not in footprint, but in *orientation*. Its
    #: thickness axis was fit by PCA over the panel's own tracked points, and
    #: on the real demo episode that axis came out **41.7 degrees off** the
    #: drawer's independently-measured prismatic slide axis, making the slab
    #: 0.415 m thick along the true slide direction (14x the intended
    #: 0.03 m) -- visible as a long diagonal bar sweeping through
    #: ``control_depth.mkv`` as the drawer opened. A *fitted* thickness axis
    #: -- PCA'd, slide-axis-constrained, or otherwise -- was the wrong kind
    #: of fix: it is still a guess standing in for a measurement that already
    #: exists. ``OBSERVED_POINTS`` guesses nothing: the geometry is S3's own
    #: observed depth (dense or background plate) unprojected within the S5
    #: object mask at a well-observed reference span, carried as a raw point
    #: cloud and rigidly transformed per frame by the object's own
    #: independently-fitted ``T_world_obj(t)`` track -- no orientation choice
    #: is made at all, so no orientation bug is possible. Valid only from
    #: the capture viewpoint (a single-view 2.5D shell, like every point
    #: cloud in this codebase -- see ``scripts/render_pointcloud_scene.py``'s
    #: own docstring) and only while the object's own rotation stays small
    #: enough that no genuinely new surface needs to appear (see
    #: ``scripts/run_datagen.py``'s own rotation-range gate) -- a lifted or
    #: rotated object (e.g. this episode's brick) still needs a real mesh
    #: (``SAM3D``), because new surfaces genuinely become visible that no
    #: single reference-frame capture ever observed.
    OBSERVED_POINTS = "observed_points"


class InteractionState(Enum):
    """What the robot is doing to the object on a given frame."""

    FREE = "free"
    PUSHED = "pushed"
    GRASPED = "grasped"


@dataclass
class ObjectCrop:
    """A tight RGBA crop of the object, ready for image-to-3D."""

    frame_idx: int
    rgba: np.ndarray  # (h, w, 4) uint8; alpha = mask
    bbox_xyxy: tuple[int, int, int, int]  # in full-frame pixels
    mask_full: np.ndarray  # (H, W) bool at full frame resolution
    prompt: str = ""
    #: Fraction of the mask that was occluded and filled in, 0 if no inpainting.
    inpainted_fraction: float = 0.0

    @property
    def area_px(self) -> int:
        return int(self.mask_full.sum())


@dataclass
class ObjectMesh:
    """A mesh in its own canonical (unit-ish, origin-centred) frame."""

    vertices: np.ndarray  # (V, 3) float64
    faces: np.ndarray  # (F, 3) int64
    source: MeshSource
    vertex_colors: np.ndarray | None = None  # (V, 3) uint8
    texture: np.ndarray | None = None  # (H, W, 3) uint8
    uv: np.ndarray | None = None  # (V, 2) float32
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def extent(self) -> np.ndarray:
        return self.vertices.max(axis=0) - self.vertices.min(axis=0)

    @property
    def centroid(self) -> np.ndarray:
        return self.vertices.mean(axis=0)


@dataclass
class ObjectAlignment:
    """Places an :class:`ObjectMesh` into the world (robot-base) frame.

    ``transform`` maps mesh-canonical coordinates to world metres, including
    ``scale``. Keeping the observed cloud alongside it is what makes the fit
    auditable: the render can always show mesh-vs-cloud.
    """

    transform: np.ndarray  # (4, 4) mesh -> world, scale baked in
    scale: float
    points_world: np.ndarray  # (N, 3) observed object points, world frame
    point_colors: np.ndarray | None = None  # (N, 3) uint8
    rmse_m: float = float("nan")
    inlier_fraction: float = float("nan")
    frame_idx: int = 0
    #: Whether ``transform``'s *rotation* was actually fit (multi-start ICP
    #: against a sign-searched PCA init) or is an unfit axis-aligned/identity
    #: fallback because the observed cloud could not support one -- see
    #: :func:`fpgm.objects.align.fit_mesh_to_points`. ``True`` even when
    #: ``refine_icp=False`` was requested (PCA-only is still a fit, just not
    #: ICP-refined); ``False`` only for the no-orientation-evidence fallback.
    orientation_confident: bool = True
    #: Where ``scale`` came from: ``"mask_angular_extent"`` (mask bbox x
    #: median depth -- the measured-accurate source for small objects),
    #: ``"depth_extent"`` (point-cloud spatial extent -- fine for objects
    #: large relative to depth noise, wrong for centimetre-scale ones), or
    #: ``"fixed"`` (``allow_scale=False``, scale pinned to 1.0).
    scale_source: str = "depth_extent"
    #: Human-readable provenance, e.g. why orientation was not trusted. Empty
    #: string if there is nothing noteworthy to report.
    notes: str = ""

    @property
    def position(self) -> np.ndarray:
        return self.transform[:3, 3]


@dataclass
class ObjectTrajectory:
    """Per-frame object pose produced by the interaction model."""

    timestamps: np.ndarray  # (T,) seconds
    transforms: np.ndarray  # (T, 4, 4) mesh -> world
    states: list[InteractionState]
    #: Distance from the gripper to the object surface each frame, metres.
    gripper_distance_m: np.ndarray | None = None
    #: Amplification applied to pushed motion. 1.0 = physical.
    push_gain: float = 1.0

    def __len__(self) -> int:
        return len(self.timestamps)

    @property
    def positions(self) -> np.ndarray:
        return self.transforms[:, :3, 3]

    def state_spans(self) -> list[tuple[InteractionState, int, int]]:
        """Contiguous ``(state, start, end)`` runs -- a readable timeline."""
        spans: list[tuple[InteractionState, int, int]] = []
        if not self.states:
            return spans
        cur, start = self.states[0], 0
        for i, s in enumerate(self.states[1:], 1):
            if s is not cur:
                spans.append((cur, start, i))
                cur, start = s, i
        spans.append((cur, start, len(self.states)))
        return spans


@dataclass
class StageArtifacts:
    """Where a stage wrote its debug output."""

    stage: str
    directory: Path
    files: dict[str, Path] = field(default_factory=dict)
    stats: dict[str, Any] = field(default_factory=dict)

    def add(self, key: str, path: Path) -> Path:
        self.files[key] = path
        return path