src/fpgm/config.py

"""Declarative configuration: dataclass schema + YAML loader.

One config object drives a whole run. Every tunable that a reviewer might want to
change lives here rather than being buried as a literal inside an algorithm.
"""

from __future__ import annotations

from dataclasses import dataclass, field, fields, is_dataclass
from pathlib import Path
from typing import Any, TypeVar, get_type_hints

import yaml

T = TypeVar("T")

REPO_ROOT = Path(__file__).resolve().parents[2]


@dataclass
class PathsConfig:
    data_dir: Path = REPO_ROOT / "data"
    checkpoints_dir: Path = REPO_ROOT / "checkpoints"
    outputs_dir: Path = REPO_ROOT / "outputs"

    @property
    def pointworld_dir(self) -> Path:
        return self.data_dir / "pointworld"

    @property
    def droid_dir(self) -> Path:
        return self.data_dir / "droid_raw"


@dataclass
class DatasetConfig:
    """Which episodes to pull, and from where."""

    hf_repo: str = "nvidia/PointWorld-DROID"
    shard: str = "shard-000000"
    episode_uuids: list[str] = field(default_factory=list)
    max_episodes: int = 5
    #: which exterior camera to run the pipeline on
    camera_role: str = "ext1"


@dataclass
class SegmentationConfig:
    #: FlashAttention-3 is SAM 3.1's default but needs a separate install; opt in.
    use_fa3: bool = False
    multiplex_count: int = 16
    max_num_objects: int = 16
    #: The frame tensor is ~6.1 MB/frame on GPU; this machine's GPUs are shared.
    offload_video_to_cpu: bool = True
    offload_state_to_cpu: bool = False
    output_prob_thresh: float = 0.5
    prompt_frame_idx: int = 0
    #: explicit per-episode prompt override, keyed by episode uuid
    prompt_overrides: dict[str, str] = field(default_factory=dict)
    #: Tiebreaker when a prompt matches several instances. Off by default: the
    #: deliverable is the object's own motion, so SAM's detection score decides.
    select_by_gripper_proximity: bool = False


@dataclass
class TrackingConfig:
    checkpoint: str = "tapnextpp_512.ckpt"
    input_resolution: int = 512
    #: number of query points sampled inside the object mask
    n_query_points: int = 64
    #: "farthest" spreads points via farthest-point sampling; "grid" uses a lattice
    sampling: str = "farthest"
    #: keep query points this many pixels away from the mask boundary
    boundary_erosion_px: int = 3
    autocast: bool = True


@dataclass
class LiftingConfig:
    """How 2D tracks become metric 3D.

    ``rigid_pnp`` needs depth at only ONE seed frame and then solves the object's
    6-DOF pose per frame from the 2D tracks alone. That is both cheaper and
    stronger than ``per_frame``: a rigid body has 6 unknowns against 2Q equations,
    whereas per-frame lifting needs a depth lookup for every point at every frame
    and discards the rigidity constraint entirely.

    Use ``per_frame`` for objects that are not rigid (cloth, liquid, articulated).
    """

    mode: str = "rigid_pnp"  # "rigid_pnp" | "per_frame"
    pnp_reprojection_threshold_px: float = 4.0
    pnp_min_points: int = 6


@dataclass
class DepthConfig:
    k_neighbors: int = 8
    max_image_radius_px: float = 40.0
    #: drop neighbours whose depth deviates more than this fraction from the median
    depth_discontinuity_ratio: float = 0.15
    min_support: int = 3
    #: confidence multiplier applied when falling back to unmasked neighbours
    unmasked_fallback_penalty: float = 0.5


@dataclass
class ConventionConfig:
    in_bounds_weight: float = 0.5
    color_weight: float = 0.5
    #: below this score gap the verdict is ambiguous and we refuse to guess
    min_margin: float = 0.15
    #: if the winner is still below this, calibration itself is suspect
    min_frac_in_bounds: float = 0.5


@dataclass
class VelocityConfig:
    method: str = "savgol"  # "savgol" | "central_diff"
    savgol_window: int = 5
    savgol_polyorder: int = 2
    min_confidence: float = 0.3
    #: occlusion gaps up to this long are bridged in position before differentiating
    max_gap_frames: int = 3
    #: MAD trimming factor for robust aggregation to a single object velocity
    mad_trim_factor: float = 3.0


@dataclass
class RuntimeConfig:
    device: str = "cuda"
    #: pinned explicitly because the GPUs on this box are shared
    gpu_index: int | None = None
    seed: int = 42
    log_level: str = "INFO"


@dataclass
class DatagenConfig:
    """Thresholds/resolutions/paths for the DROID -> VACE datagen pipeline (S1-S8).

    Master data (poses/masks/tracks) is kept resolution-free -- only the S8 export
    tier is rendered at a fixed size -- so ``export_resolution`` and
    ``native_resolution`` are two different numbers, not one: re-exporting at
    720p later must not require re-running S2-S7.
    """

    #: root directory datagen writes ``<uuid>/<camera>/{master,vace}/`` under.
    output_root: Path = REPO_ROOT / "outputs" / "datagen"
    #: (width, height) every S8 ``vace/`` buffer is rendered/encoded at.
    export_resolution: tuple[int, int] = (832, 480)
    #: (width, height) S3's dense-depth working grid -- matches PointWorld's own
    #: initial_depth anchor resolution, so the anchor needs no rescale to merge.
    native_resolution: tuple[int, int] = (320, 180)
    #: S3: reject splatted depth below this range as stereo-saturation noise, not
    #: real close-range geometry (background-only floor; grasped objects can be
    #: legitimately closer and are excluded from this check).
    depth_min_floor_m: float = 0.40
    #: S3: number of push-pull pyramid levels used to fill unsupported pixels.
    push_pull_levels: int = 7
    #: DEPRECATED as a gate -- kept only as the threshold for a reported
    #: diagnostic. Measured on the demo episode: the render legitimately
    #: includes the Robotiq gripper (needed downstream) while SAM3's "robot
    #: arm" reference mask mostly does not, so raw IoU penalizes a *correct*
    #: render for being more complete than the text-prompted reference --
    #: precision collapses (e.g. 0.078 on a frame where recall was still
    #: 0.724) even when the geometry is exactly right. See robot_recall_gate
    #: and robot_shift_gate_px, which replace this for gating.
    robot_iou_gate: float = 0.5
    #: S2 gate: minimum median recall (fraction of SAM3's "robot arm" mask
    #: covered by the rendered mask) across every video frame SAM3 tracked.
    #: Over-coverage (gripper, cables) is harmless and expected -- MuJoCo
    #: will produce it too at inference -- so only under-coverage is
    #: penalized here, unlike IoU/precision.
    robot_recall_gate: float = 0.85
    #: S2 gate: maximum *systematic* (mean, not median) best-fit pixel shift
    #: (each frame's shift searched over +/-24 px) between the rendered mask
    #: and SAM3's mask, across every tracked frame. A genuine calibration
    #: error is a *bias* term -- present in every frame -- so it is the
    #: *mean* of the signed per-frame (dy, dx) that isolates it: mean cancels
    #: zero-centred per-frame scatter (shape/occlusion/gripper disagreement,
    #: SAM tracking noise) the same way it would cancel measurement noise
    #: around a true value, and does not cancel a bias that is the same sign
    #: every frame. An earlier version of this gate used
    #: ``median(|shift|)`` instead, which is wrong for this purpose: a
    #: magnitude is non-negative by construction, so its median is inflated
    #: by scatter alone even at zero true bias, and it cannot distinguish a
    #: real 2px calibration error from 10px of direction-varying noise --
    #: measured on the demo episode, ``median(|shift|)`` read ~4.2px and
    #: ~7.0px for the two extrinsics candidates (both "fail", no discriminating
    #: signal), while the corrected mean-of-components bias read 1.24px and
    #: 6.20px respectively -- a 5x separation the old statistic could not see.
    #: This is not a relaxed threshold to force a pass: the corrected metric
    #: still fails the wrong candidate by a wide margin.
    robot_systematic_shift_gate_px: float = 2.0
    #: DEPRECATED for gating purposes -- kept, not removed, only because
    #: fpgm.datagen.types.DenseDepthResult's docstring (sibling-owned, out of
    #: scope to edit here) still names this field. Per-frame scene_flows
    #: splat coverage (what this originally gated) was measured to be
    #: structurally ~35% regardless of episode/camera quality -- it cannot
    #: discriminate a good capture from a bad one, so it is no longer used as
    #: a pass/fail gate. See background_coverage_gate, which replaces it.
    depth_support_gate: float = 0.6
    #: S3/S7 gate: minimum pre-fill coverage of the *background* depth plate
    #: (fpgm.datagen.background_depth.BackgroundDepthResult, built from the
    #: dense per-clip initial_depth anchors, not the sparse per-frame splat).
    #: Unlike depth_support_gate this actually varies with capture quality --
    #: measured 0.913 (ext1) / 0.956 (ext2) on the demo episode -- so 0.85
    #: discriminates a usable capture without being toothless.
    background_coverage_gate: float = 0.85
    #: S4: TAPNext full-frame query grid, (n_cols, n_rows).
    tapnext_grid: tuple[int, int] = (40, 24)
    #: S4: DBSCAN clustering of (3D position, displacement vector) tracks into
    #: candidate dynamic objects.
    dbscan_eps_m: float = 0.05
    dbscan_min_samples: int = 8
    #: S6 gate: PnP-RANSAC median reprojection error, in pixels.
    pnp_reproj_gate_px: float = 3.0
    #: S7 gate: PrismaticFit residual, in metres, below which a track counts as
    #: prismatic (e.g. the drawer front) rather than a general rigid motion.
    prismatic_residual_gate_m: float = 0.01
    #: S7: minimum ``ObjectTrack.visible_frac`` a frame must have to be kept
    #: in the *gated* PrismaticFit (see fpgm.datagen.events.PrismaticFitter's
    #: docstring). Measured on the demo episode's drawer track: gating the
    #: 127-frame track at this threshold keeps 100 frames, drops the residual
    #: 5.18mm -> 0.38mm (13x), and rotation range 14.13deg -> 1.01deg, while
    #: visible_frac < 0.5 holds for exactly the low-point-support tail
    #: (frames 100-126) -- not scattered across the track, which is what
    #: makes the tail's failure an evidence problem rather than real motion.
    prismatic_visible_frac_gate: float = 0.5
    #: S6/S7: gap runs (GAP pose_source) longer than this many frames are flagged
    #: rather than bridged -- a long gap is exactly where the object may have
    #: changed direction, the same reasoning as VelocityConfig.max_gap_frames.
    max_gap_run_frames: int = 15
    #: S8: constant-rate-factor for the real-video target.mp4 (lower = higher
    #: quality/larger file).
    target_video_crf: int = 16
    #: S8: crf for the lossy control_normal.mp4/fg_mask.mp4 renders. The other
    #: control buffers (control_depth.mkv, control_seg.mkv) are FFV1-lossless and
    #: have no crf.
    control_video_crf: int = 18
    #: S8: dilation (pixels) applied to the rendered robot U objects mask before
    #: writing fg_mask.mp4 -- a tight mask clips true object boundaries under
    #: VACE's own downstream augmentation.
    fg_mask_dilation_px: int = 4
    #: S8: frames per exported window. Wan2.1-VACE's own ``--src_video`` path
    #: ignores ``--frame_num`` and instead derives the output frame count as
    #: ``4n+1``, capped at 81 for the >=480P bucket (see
    #: fpgm.datagen.export_vace's module docstring for the pinned-source
    #: citation) -- 81 = 4*20+1 is the largest value satisfying that contract,
    #: so an episode longer than 81 frames must be split into overlapping
    #: windows rather than exported as one clip.
    vace_window_frames: int = 81
    #: S8: stride (frames) between successive window starts. Chosen as
    #: exactly half of vace_window_frames (40, ~49% overlap) -- enough
    #: overlap that an event straddling a window boundary (e.g. the drawer
    #: closing) still lands fully inside at least one window, while still
    #: roughly halving the number of windows a naive stride=1 would produce.
    #: Not a claim that 40 is uniquely optimal, just a documented, explicit
    #: choice per the task brief ("pick and document a stride") -- see
    #: fpgm.datagen.export_vace.compute_windows for how it is actually used
    #: (the final window is always anchored to end exactly at n_frames, so
    #: this stride governs the *interior* windows, not the last one).
    vace_window_stride: int = 40
    #: S8 gate, from the pinned Wan2.1 source (see export_vace.py's module
    #: docstring): output resolution derives from the control video's own
    #: pixel size, and needs ``(h//16)*(w//16) >= 1560`` to reach the
    #: documented 81-frame/832x480 bucket -- a smaller control silently
    #: yields a different, smaller output resolution with no error. Checked
    #: against ``export_resolution`` at export time, not just documented.
    vace_min_hw16_product: int = 1560

    def __post_init__(self) -> None:
        # YAML round-trips these as lists; normalise to tuples so callers can
        # rely on the annotated type (and so two configs loaded differently
        # compare equal).
        self.export_resolution = tuple(self.export_resolution)  # type: ignore[assignment]
        self.native_resolution = tuple(self.native_resolution)  # type: ignore[assignment]
        self.tapnext_grid = tuple(self.tapnext_grid)  # type: ignore[assignment]


@dataclass
class PipelineConfig:
    paths: PathsConfig = field(default_factory=PathsConfig)
    dataset: DatasetConfig = field(default_factory=DatasetConfig)
    segmentation: SegmentationConfig = field(default_factory=SegmentationConfig)
    tracking: TrackingConfig = field(default_factory=TrackingConfig)
    depth: DepthConfig = field(default_factory=DepthConfig)
    lifting: LiftingConfig = field(default_factory=LiftingConfig)
    convention: ConventionConfig = field(default_factory=ConventionConfig)
    velocity: VelocityConfig = field(default_factory=VelocityConfig)
    runtime: RuntimeConfig = field(default_factory=RuntimeConfig)
    datagen: DatagenConfig = field(default_factory=DatagenConfig)

    @classmethod
    def from_yaml(cls, path: str | Path) -> "PipelineConfig":
        with open(path) as fh:
            raw = yaml.safe_load(fh) or {}
        return _build(cls, raw)

    def to_yaml(self, path: str | Path) -> None:
        with open(path, "w") as fh:
            yaml.safe_dump(_unbuild(self), fh, sort_keys=False)


def _build(cls: type[T], raw: dict[str, Any]) -> T:
    """Recursively instantiate nested dataclasses, rejecting unknown keys.

    Unknown keys are an error rather than being ignored: a silently-dropped typo in
    a config is indistinguishable from the setting having no effect.

    Field types are resolved with ``get_type_hints`` rather than read off
    ``Field.type``: this module uses ``from __future__ import annotations``, so the
    raw annotations are *strings* and a nested dataclass would otherwise go
    undetected and be assigned as a plain dict.
    """
    known = {f.name for f in fields(cls)}  # type: ignore[arg-type]
    unknown = set(raw) - known
    if unknown:
        raise ValueError(f"unknown config keys for {cls.__name__}: {sorted(unknown)}")

    hints = get_type_hints(cls)
    kwargs: dict[str, Any] = {}
    for name, value in raw.items():
        ftype = hints.get(name)
        if is_dataclass(ftype) and isinstance(value, dict):
            kwargs[name] = _build(ftype, value)  # type: ignore[arg-type]
        elif ftype is Path and value is not None:
            kwargs[name] = Path(value)
        else:
            kwargs[name] = value
    return cls(**kwargs)  # type: ignore[return-value]


def _unbuild(obj: Any) -> Any:
    if is_dataclass(obj):
        return {f.name: _unbuild(getattr(obj, f.name)) for f in fields(obj)}
    if isinstance(obj, Path):
        return str(obj)
    if isinstance(obj, dict):
        return {k: _unbuild(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_unbuild(v) for v in obj]
    return obj