src/fpgm/geometry/lifting.py

"""Lift a 2D point track to metric 3D world positions via a depth source.

Kept separate from :mod:`fpgm.geometry.camera` and :mod:`fpgm.depth.scene_flow`
because it is the glue step: it owns none of the projection math or interpolation
algorithm, only the per-frame loop that ties a track, a camera, a depth source, and
timing together into a :class:`~fpgm.types.Track3D`.
"""

from __future__ import annotations

import numpy as np

from fpgm.depth.base import DepthSource
from fpgm.geometry.camera import Camera
from fpgm.types import ClipTiming, DepthQuery, NoValidDepthAnnotationsError, Track2D, Track3D


def lift_track_to_3d(
    track2d: Track2D,
    camera: Camera,
    depth_source: DepthSource,
    timing: ClipTiming,
    object_masks: dict[int, np.ndarray] | None = None,
) -> Track3D:
    """Lift a 2D track to a :class:`~fpgm.types.Track3D` of world-frame positions.

    Loops frames, queries ``depth_source`` for each frame's visible query points,
    unprojects the resulting ``(u, v, depth)`` triples to world coordinates, and
    fills in confidence/validity from the depth query's own provenance.

    Note:
        ``track2d.frames`` must already be clip-local integer frame indices aligned
        with the ``scene_flows`` T axis -- i.e. ``track2d`` should already have been
        passed through
        :func:`fpgm.geometry.temporal.resample_track_to_clip_frames` if it was
        originally mp4-frame indexed. ``timing`` is used only to convert those
        clip-local indices to wall-clock timestamps.

    Args:
        track2d: 2D track, uv valid at ``track2d.resolution``.
        camera: :class:`~fpgm.geometry.camera.Camera` at the scene-flow annotation
            resolution; rescaled internally to ``track2d.resolution`` if they differ
            -- this is the explicit, single place that resolution mismatch is handled.
        depth_source: Source of per-frame metric depth.
        timing: Used to convert clip-local frame indices to timestamps.
        object_masks: Optional ``{clip_frame_idx: (H, W) bool}`` masks, at
            ``track2d.resolution``, restricting depth interpolation support per frame.

    Returns:
        A :class:`~fpgm.types.Track3D` with the same point identities as ``track2d``.
    """
    width, height = track2d.resolution
    cam = camera.rescaled(width, height)

    n_frames = track2d.frames.shape[0]
    n_points = track2d.uv.shape[1]

    timestamps = np.asarray(timing.clip_frame_to_seconds(track2d.frames.astype(np.int64)))
    xyz_world = np.full((n_frames, n_points, 3), np.nan, dtype=np.float64)
    valid = np.zeros((n_frames, n_points), dtype=bool)
    confidence = np.zeros((n_frames, n_points), dtype=np.float32)

    for i, frame_idx in enumerate(track2d.frames):
        frame_idx = int(frame_idx)
        vis = track2d.visible[i]
        if not np.any(vis):
            continue
        point_indices = np.flatnonzero(vis)
        uv = track2d.uv[i, point_indices]

        mask = object_masks.get(frame_idx) if object_masks is not None else None
        query = DepthQuery(
            frame_idx=frame_idx, uv=uv, query_resolution=(width, height), object_mask=mask
        )
        try:
            result = depth_source.query(query)
        except NoValidDepthAnnotationsError:
            # No usable scene-flow support this frame at all -- leave every point
            # for this frame invalid rather than raising, so a single bad frame
            # does not abort lifting the whole track.
            continue

        good = point_indices[result.valid]
        good_uv = track2d.uv[i, good]
        good_depth = result.depth[result.valid]
        world_xyz = cam.unproject(good_uv, good_depth)

        xyz_world[i, good] = world_xyz
        valid[i, good] = True
        confidence[i, good] = result.confidence[result.valid]

    return Track3D(
        point_id=track2d.point_id,
        timestamps=timestamps,
        xyz_world=xyz_world,
        valid=valid,
        confidence=confidence,
    )