src/fpgm/types.py

"""Shared data contracts for the whole pipeline.

Every stage consumes and produces the dataclasses defined here. No stage reaches
into another stage's internal state: the segmenter never sees a checkpoint path,
the velocity estimator never sees an HDF5 handle. This keeps each stage
independently unit-testable against synthetic data.

Array shape conventions used throughout:
    T = frames in a clip, N = annotated scene-flow points, Q = tracked query points

Pixel coordinates are always ``(u, v)`` = ``(x, y)`` and are always paired with the
resolution they are valid at -- see :class:`CameraIntrinsics`.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any

import numpy as np

# --------------------------------------------------------------------------- #
# Exceptions
# --------------------------------------------------------------------------- #


class FpgmError(Exception):
    """Base class for every error raised by this package."""


class ConfigError(FpgmError):
    """A declared configuration is internally inconsistent or incomplete.

    Distinct from :class:`DataError`: nothing was wrong with the episode, the
    *configuration* names something that cannot be satisfied (a mesh role with
    no registered mesh, a parent label naming an object this episode does not
    have). Raised at resolve time so the failure lands next to the declaration
    rather than deep inside a stage.
    """


class GeometryError(FpgmError):
    """Base class for geometry/lifting failures."""


class AmbiguousConventionError(GeometryError):
    """Scene-flow frame convention could not be decided with a decisive margin.

    Raised rather than defaulting, because guessing wrong applies a systematic
    rigid transform to every 3D position while still looking plausible.
    """


class CameraCalibrationMismatchError(GeometryError):
    """Both frame hypotheses score badly -- intrinsic/extrinsic/serial mispairing."""


class NoValidDepthAnnotationsError(GeometryError):
    """A frame has no usable scene-flow support at all.

    Distinct from "this point was rejected" so downstream gap handling can tell
    "no data this frame" apart from "this individual point is an outlier".
    """


class InsufficientTrackHistoryError(GeometryError):
    """Not enough valid samples to differentiate."""


class DataError(FpgmError):
    """Base class for data acquisition / parsing failures."""


class EpisodeNotFoundError(DataError):
    """Episode could not be located in the DROID raw bucket."""


# --------------------------------------------------------------------------- #
# Episode identity
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class EpisodeId:
    """A DROID episode uuid, e.g. ``AUTOLab+0d4edc83+2023-10-21-19h-07m-04s``.

    The uuid deterministically encodes the episode's location in the public
    DROID raw bucket; see :mod:`fpgm.data.ids` for the mapping.
    """

    uuid: str

    @property
    def lab(self) -> str:
        return self.uuid.split("+")[0]

    @property
    def user_id(self) -> str:
        return self.uuid.split("+")[1]

    @property
    def timestamp(self) -> datetime:
        return datetime.strptime(self.uuid.split("+")[2], "%Y-%m-%d-%Hh-%Mm-%Ss")

    def __str__(self) -> str:  # pragma: no cover - trivial
        return self.uuid


# --------------------------------------------------------------------------- #
# Camera
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class CameraIntrinsics:
    """Pinhole intrinsics that always carry the resolution they are valid at.

    Carrying ``(width, height)`` is what makes resolution mismatches impossible to
    introduce silently: the scene-flow ``intrinsic`` in the HDF5 is valid at the
    annotation resolution, while mp4 frames are typically a different size.
    """

    fx: float
    fy: float
    cx: float
    cy: float
    width: int
    height: int

    def scaled(self, width: int, height: int) -> "CameraIntrinsics":
        """Return intrinsics rescaled to a different pixel resolution."""
        sx = width / self.width
        sy = height / self.height
        return CameraIntrinsics(
            fx=self.fx * sx,
            fy=self.fy * sy,
            cx=self.cx * sx,
            cy=self.cy * sy,
            width=width,
            height=height,
        )

    @property
    def matrix(self) -> np.ndarray:
        """3x3 K matrix (no skew)."""
        return np.array(
            [[self.fx, 0.0, self.cx], [0.0, self.fy, self.cy], [0.0, 0.0, 1.0]],
            dtype=np.float64,
        )

    @classmethod
    def from_matrix(cls, k: np.ndarray, width: int, height: int) -> "CameraIntrinsics":
        k = np.asarray(k, dtype=np.float64)
        return cls(
            fx=float(k[0, 0]),
            fy=float(k[1, 1]),
            cx=float(k[0, 2]),
            cy=float(k[1, 2]),
            width=int(width),
            height=int(height),
        )


class FrameConvention(Enum):
    """Which frame ``scene_flows`` positions are expressed in."""

    WORLD = "world"
    CAMERA = "camera"


@dataclass(frozen=True)
class ConventionDetectionResult:
    convention: FrameConvention
    score_world: float
    score_camera: float
    margin: float
    n_points_scored: int
    frac_in_bounds: float


# --------------------------------------------------------------------------- #
# Depth
# --------------------------------------------------------------------------- #


class DepthMethod(Enum):
    """Why a depth value is (or is not) what it is.

    Carrying a reason code rather than a bare NaN is deliberate: debugging "why is
    this point's velocity NaN" three stages downstream without provenance is the
    single biggest time sink in a pipeline like this.
    """

    DIRECT = "direct"
    INTERP = "interp"
    INTERP_UNMASKED_FALLBACK = "interp_unmasked_fallback"
    REJECTED_NO_SUPPORT = "rejected_no_support"
    REJECTED_DISCONTINUITY = "rejected_discontinuity"


@dataclass
class DepthQuery:
    """Ask a :class:`~fpgm.depth.base.DepthSource` for depth at 2D locations."""

    frame_idx: int
    uv: np.ndarray  # (Q, 2) float, pixels at ``query_resolution``
    query_resolution: tuple[int, int]  # (width, height)
    object_mask: np.ndarray | None = None  # (H, W) bool, restricts neighbour search


@dataclass
class DepthResult:
    depth: np.ndarray  # (Q,) float32, metres; NaN where invalid
    valid: np.ndarray  # (Q,) bool
    confidence: np.ndarray  # (Q,) float32 in [0, 1]
    method: np.ndarray  # (Q,) object array of DepthMethod
    n_support: np.ndarray  # (Q,) int32


# --------------------------------------------------------------------------- #
# Segmentation
# --------------------------------------------------------------------------- #


@dataclass
class FrameMasks:
    """SAM 3.1 output for a single frame, at the video's original resolution."""

    frame_idx: int
    obj_ids: np.ndarray  # (M,) int64
    masks: np.ndarray  # (M, H, W) bool
    scores: np.ndarray  # (M,) float32
    boxes_xywh: np.ndarray  # (M, 4) float32, normalised

    def mask_for(self, obj_id: int) -> np.ndarray | None:
        idx = np.flatnonzero(self.obj_ids == obj_id)
        return self.masks[int(idx[0])] if idx.size else None


@dataclass
class Masklet:
    """One object's mask across a whole clip."""

    obj_id: int
    frames: dict[int, np.ndarray] = field(default_factory=dict)  # frame_idx -> (H,W) bool
    scores: dict[int, float] = field(default_factory=dict)

    def area(self, frame_idx: int) -> int:
        m = self.frames.get(frame_idx)
        return int(m.sum()) if m is not None else 0


# --------------------------------------------------------------------------- #
# Tracking
# --------------------------------------------------------------------------- #


@dataclass
class Track2D:
    """TAPNext++ output: Q points followed across T frames, in display pixels."""

    point_id: np.ndarray  # (Q,) int32
    frames: np.ndarray  # (T,) int32, absolute video frame indices
    uv: np.ndarray  # (T, Q, 2) float32, pixels
    visible: np.ndarray  # (T, Q) bool
    resolution: tuple[int, int]  # (width, height) the uv values are valid at


@dataclass
class Track3D:
    """Lifted metric 3D positions in the world (robot-base) frame."""

    point_id: np.ndarray  # (Q,) int32
    timestamps: np.ndarray  # (T,) float64, seconds
    xyz_world: np.ndarray  # (T, Q, 3) float64, metres; NaN where invalid
    valid: np.ndarray  # (T, Q) bool
    confidence: np.ndarray  # (T, Q) float32


@dataclass
class VelocityEstimate:
    """Per-point and aggregated object velocity, in the world frame, in m/s."""

    point_id: np.ndarray  # (Q,)
    timestamps: np.ndarray  # (T,) float64, seconds
    linear_velocity: np.ndarray  # (T, Q, 3) m/s, NaN where unavailable
    object_linear_velocity: np.ndarray  # (T, 3) m/s
    object_speed: np.ndarray  # (T,) m/s
    object_velocity_confidence: np.ndarray  # (T,)
    inlier_mask: np.ndarray  # (T, Q) bool
    frame: str = "world"  # reference-frame provenance, carried explicitly


# --------------------------------------------------------------------------- #
# Timing
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class ClipTiming:
    """Reconciles the three clocks: clip index, video frame index, wall-clock seconds.

    ``scene_flows`` shares its T axis with ``joint_positions``/``gripper_pose``, so
    ``trajectory_fps`` (DROID's control rate, 15 Hz) is the canonical clock and is
    the only thing used to turn indices into seconds.

    Two independent time-base traps have been found here, and both matter:

    1. **The mp4-fps trap.** Index alignment is derived from **frame counts**, never
       from ``mp4_fps``. On real DROID recordings the container header reports 60 fps
       while the file holds 127 frames for a 128-step trajectory -- i.e. the true
       correspondence is one video frame per trajectory step. Trusting the header
       would stretch the time axis by 4x and scale every velocity by the same
       factor, which is exactly the kind of error that produces plausible-looking
       but wrong numbers. ``mp4_fps`` is therefore kept for diagnostics and playback
       only; ``video_frames_per_step`` (measured from frame counts) is what is
       actually used.

    2. **The annotation-stride trap.** A PointWorld scene-flow clip's own frame axis
       (what ``clip_start_frame``/``clip_end_frame``/clip-local ``t`` are expressed
       in) is *not* one-to-one with the trajectory/video frame axis either. Clip
       frame ``t`` of clip ``"50:61"`` is trajectory/video row ``2 * (50 + t)``, not
       ``50 + t`` -- the annotations (``scene_flows``, ``joint_positions``,
       ``gripper_pose``, all sharing the clip group's T axis) are recorded at half
       the trajectory/video rate. This was measured, not assumed: matching each
       clip's own ``joint_positions`` against ``trajectory.h5``'s
       ``observation/robot_state/joint_positions`` at candidate strides gives an
       exact (0.0 max abs error) match at stride 2 across all 42 clips checked, and
       a large mismatch at stride 1. See
       :func:`fpgm.data.pointworld.verify_annotation_stride`, which performs exactly
       this check on real data -- callers should measure ``annotation_stride``
       from their own clip rather than trusting the default, for the same reason
       trap 1 is not read from the mp4 header: a wrong-but-plausible constant is
       how both of these traps happened in the first place.

    The two traps are independent and compose multiplicatively: a clip frame maps
    to a trajectory row via ``annotation_stride`` alone, and a trajectory row maps
    to a video frame via ``video_frames_per_step`` alone.
    """

    clip_start_frame: int
    clip_end_frame: int
    trajectory_fps: float
    #: Container header value. Diagnostic/playback only -- never used for alignment.
    mp4_fps: float = 0.0
    #: Video frames per trajectory step, measured from frame counts (normally 1.0).
    video_frames_per_step: float = 1.0
    #: Annotated clip frames per trajectory/video frame. Measured, not assumed --
    #: see :func:`fpgm.data.pointworld.verify_annotation_stride`.
    annotation_stride: int = 2

    @classmethod
    def from_counts(
        cls,
        clip_start: int,
        clip_end: int,
        trajectory_length: int,
        video_frame_count: int,
        trajectory_fps: float = 15.0,
        mp4_fps: float = 0.0,
        annotation_stride: int = 2,
    ) -> "ClipTiming":
        """Build timing by measuring the video/trajectory frame-count ratio.

        ``annotation_stride`` is not derived here -- it can only be measured against
        real joint-position data (see
        :func:`fpgm.data.pointworld.verify_annotation_stride`), which this
        classmethod does not have access to. Callers that have verified it should
        pass it explicitly; the default is the stride measured across every
        available clip at the time this trap was found, not a guess.
        """
        ratio = 1.0
        if trajectory_length > 0 and video_frame_count > 0:
            ratio = video_frame_count / trajectory_length
        return cls(
            clip_start_frame=int(clip_start),
            clip_end_frame=int(clip_end),
            trajectory_fps=float(trajectory_fps),
            mp4_fps=float(mp4_fps),
            video_frames_per_step=float(ratio),
            annotation_stride=int(annotation_stride),
        )

    @property
    def n_frames(self) -> int:
        return self.clip_end_frame - self.clip_start_frame

    @property
    def dt(self) -> float:
        """Seconds between consecutive clip frames.

        ``annotation_stride`` trajectory steps elapse between consecutive
        *annotated* clip frames, so this is ``annotation_stride / trajectory_fps``
        -- 2/15 s, not 1/15 s, once the stride has been measured as 2.
        """
        return self.annotation_stride / self.trajectory_fps

    def clip_frame_to_trajectory_frame(self, t: int | np.ndarray):
        """Trajectory.h5 row (``joint_positions``/``gripper_pose`` index) for clip frame ``t``."""
        return (self.clip_start_frame + t) * self.annotation_stride

    def clip_frame_to_video_frame(self, t: int | np.ndarray):
        """Video frame for clip frame ``t``, composing both time-base traps."""
        return self.clip_frame_to_trajectory_frame(t) * self.video_frames_per_step

    def clip_frame_to_seconds(self, t: int | np.ndarray):
        return self.clip_frame_to_trajectory_frame(t) / self.trajectory_fps

    def video_frame_to_seconds(self, f: int | np.ndarray):
        return f / (self.trajectory_fps * self.video_frames_per_step)

    def video_frame_to_clip_frame(self, f: int | np.ndarray):
        return f / (self.video_frames_per_step * self.annotation_stride) - self.clip_start_frame

    def timestamps(self) -> np.ndarray:
        """Wall-clock seconds for each clip frame."""
        return self.clip_frame_to_seconds(np.arange(self.n_frames, dtype=np.float64))


# --------------------------------------------------------------------------- #
# Episode / clip
# --------------------------------------------------------------------------- #


@dataclass
class CameraAssets:
    """Everything known about one physical camera for one clip."""

    serial: str
    role: str  # "ext1" | "ext2" | "wrist"
    intrinsic: np.ndarray  # (3, 3) float64, at annotation resolution
    extrinsic: np.ndarray  # (4, 4) float64, world -> camera
    mp4_path: str | None = None


@dataclass
class SceneFlowClip:
    """One ``"{start}:{end}"`` clip group of one camera in a ``*_flows.h5``."""

    key: str
    start: int
    end: int
    camera_serial: str
    scene_flows: np.ndarray  # (T, N, 3) float32
    scene_visibility: np.ndarray  # (T, N) bool
    scene_depth_valid: np.ndarray  # (T, N) bool
    scene_colors: np.ndarray  # (T, N, 3) uint8
    intrinsic: np.ndarray  # (3, 3) float64
    extrinsic: np.ndarray  # (4, 4) float64, world -> camera
    initial_rgb: np.ndarray | None = None  # (H, W, 3) uint8, decoded
    gripper_pose: np.ndarray | None = None  # (T, 7) float32
    gripper_open: np.ndarray | None = None  # (T, 1) bool
    #: (T, 7) float32, robot joint positions -- shares the clip group's T axis with
    #: gripper_pose/scene_flows. Exists mainly so the clip -> trajectory.h5 stride
    #: can be *measured* against it; see fpgm.data.pointworld.verify_annotation_stride.
    joint_positions: np.ndarray | None = None
    #: (180, 320) uint16 millimetres, dense depth for the clip's *first* frame only
    #: (there is no dense depth for the other T-1 frames -- that is the whole reason
    #: the datagen plan splats ``scene_flows`` for the rest and treats this as an
    #: anchor/scale reference rather than a full depth stream). Optional with a
    #: ``None`` default so existing callers/pickled instances built before this
    #: field existed keep working unchanged.
    initial_depth: np.ndarray | None = None
    #: (T, N, 3) int8, per-point surface normals sharing scene_flows' (T, N) axes.
    #: Signed unit-ish components stored as int8 (i.e. roughly [-127, 127] mapping
    #: to [-1, 1]) -- decoding to float is left to the caller since different
    #: consumers want different precision/normalization. Optional for the same
    #: backward-compatibility reason as initial_depth.
    scene_normals: np.ndarray | None = None

    @property
    def n_frames(self) -> int:
        return int(self.scene_flows.shape[0])

    @property
    def n_points(self) -> int:
        return int(self.scene_flows.shape[1])


@dataclass
class Episode:
    """A DROID episode joined across both data sources."""

    episode_id: EpisodeId
    task_instruction: str
    trajectory_length: int
    cameras: dict[str, CameraAssets] = field(default_factory=dict)  # serial -> assets
    metadata: dict[str, Any] = field(default_factory=dict)
    clips: list[SceneFlowClip] = field(default_factory=list)

    def camera_by_role(self, role: str) -> CameraAssets | None:
        for cam in self.cameras.values():
            if cam.role == role:
                return cam
        return None