src/fpgm/geometry/temporal.py

"""Reconcile clip-key parsing and track resampling against :class:`~fpgm.types.ClipTiming`.

Three clocks are in play: the ``"{start}:{end}"`` clip key (trajectory-fps frame
indices), the mp4 container's own frame indices (``mp4_fps``), and wall-clock
seconds. :class:`~fpgm.types.ClipTiming` exposes the conversions; this module builds
one from a clip key and uses it to resample a :class:`~fpgm.types.Track2D` (mp4-frame
indexed) onto the integer clip-frame axis that ``scene_flows`` shares.
"""

from __future__ import annotations

import numpy as np

from fpgm.types import ClipTiming, Track2D


def build_timing(clip_key: str, trajectory_fps: float, mp4_fps: float) -> ClipTiming:
    """Parse a ``"{start}:{end}"`` scene-flow clip group key into a :class:`ClipTiming`.

    Args:
        clip_key: e.g. ``"1204:1450"``, the h5 group name for one clip.
        trajectory_fps: fps of the ``scene_flows``/``joint_positions`` clock.
        mp4_fps: fps read from the mp4 container (not assumed equal to ``trajectory_fps``).

    Raises:
        ValueError: if ``clip_key`` is not of the form ``"{int}:{int}"``.
    """
    parts = clip_key.split(":")
    if len(parts) != 2:
        raise ValueError(f"clip key {clip_key!r} is not of the form '{{start}}:{{end}}'")
    try:
        start, end = int(parts[0]), int(parts[1])
    except ValueError as exc:
        raise ValueError(f"clip key {clip_key!r} has non-integer bounds") from exc
    if end <= start:
        raise ValueError(f"clip key {clip_key!r} has end <= start")
    return ClipTiming(
        clip_start_frame=start, clip_end_frame=end, trajectory_fps=trajectory_fps, mp4_fps=mp4_fps
    )


def resample_track_to_clip_frames(track2d: Track2D, timing: ClipTiming) -> Track2D:
    """Resample an mp4-frame-indexed :class:`Track2D` onto integer clip-frame indices.

    ``track2d.frames`` holds absolute mp4 frame indices; the output is re-indexed to
    the ``[0, timing.n_frames)`` clip-local frame axis that ``scene_flows`` shares,
    with ``uv`` **linearly** interpolated in time (not nearest-frame snapped, since
    DROID's mp4 and trajectory clocks are rarely frame-aligned). Clip frames whose
    corresponding continuous mp4-time position falls outside ``track2d.frames``'
    covered range are marked ``visible=False`` rather than extrapolated -- an
    extrapolated position outside the tracked window is fabricated data.

    Note:
        The returned ``Track2D.frames`` holds clip-local integer frame indices
        (0-based, aligned with the ``scene_flows`` T axis), not mp4 frame indices --
        this is what downstream depth queries and lifting index by.

    Args:
        track2d: Source track, uv valid at ``track2d.resolution``, indexed by
            ``track2d.frames`` (absolute mp4 frame indices).
        timing: Clock reconciliation for the clip being resampled onto.

    Returns:
        A new :class:`Track2D` with ``frames = arange(timing.n_frames)``.
    """
    n_out = timing.n_frames
    clip_frames = np.arange(n_out, dtype=np.int64)
    seconds = timing.clip_frame_to_seconds(clip_frames)
    query_video_frame = seconds * timing.mp4_fps  # (n_out,) float positions

    src_frames = track2d.frames.astype(np.float64)  # (T_src,)
    n_src = src_frames.shape[0]
    n_points = track2d.uv.shape[1]

    if n_src < 2:
        # A single source sample cannot be linearly interpolated; only an exact
        # time match is meaningful, everything else is out of coverage.
        out_uv = np.zeros((n_out, n_points, 2), dtype=np.float32)
        out_visible = np.zeros((n_out, n_points), dtype=bool)
        if n_src == 1:
            exact = np.isclose(query_video_frame, src_frames[0])
            out_uv[exact] = track2d.uv[0]
            out_visible[exact] = track2d.visible[0]
        return Track2D(
            point_id=track2d.point_id,
            frames=clip_frames.astype(np.int32),
            uv=out_uv,
            visible=out_visible,
            resolution=track2d.resolution,
        )

    in_coverage = (query_video_frame >= src_frames[0]) & (query_video_frame <= src_frames[-1])

    # Vectorised (over both output-time and query points at once) linear
    # interpolation via searchsorted -- no Python loop over points.
    right_idx = np.searchsorted(src_frames, query_video_frame, side="right")
    idx1 = np.clip(right_idx, 1, n_src - 1)
    idx0 = idx1 - 1

    t0 = src_frames[idx0]
    t1 = src_frames[idx1]
    denom = t1 - t0
    frac = np.where(denom > 0, (query_video_frame - t0) / np.where(denom > 0, denom, 1.0), 0.0)
    frac = frac.reshape(n_out, 1, 1)

    uv0 = track2d.uv[idx0]  # (n_out, Q, 2)
    uv1 = track2d.uv[idx1]  # (n_out, Q, 2)
    out_uv = (uv0 * (1.0 - frac) + uv1 * frac).astype(np.float32)

    vis0 = track2d.visible[idx0]  # (n_out, Q)
    vis1 = track2d.visible[idx1]  # (n_out, Q)
    out_visible = vis0 & vis1 & in_coverage.reshape(n_out, 1)

    out_uv = np.where(out_visible[..., None], out_uv, np.nan).astype(np.float32)

    return Track2D(
        point_id=track2d.point_id,
        frames=clip_frames.astype(np.int32),
        uv=out_uv,
        visible=out_visible,
        resolution=track2d.resolution,
    )