src/fpgm/tracking/tapnext.py

"""TAPNext++ point tracker adapter.

TAPNext++ (``google-deepmind/tapnet``) is a strictly causal, O(1)-memory-per-frame
point tracker driven through a recurrent state: ``track_frame`` either seeds new
queries (``state=None``) or advances an existing rollout (``state=<previous
state>``). The one hard constraint that shapes this whole module is that queries
can only be seeded once, on a single frame -- see :class:`TapNextPointTracker`.
"""

from __future__ import annotations

import logging
from collections.abc import Iterable, Iterator

import numpy as np

from fpgm.config import TrackingConfig
from fpgm.tracking.base import PointTracker
from fpgm.types import Track2D

logger = logging.getLogger(__name__)


class TapNextPointTracker(PointTracker):
    """Adapts ``tapnet.tapnextpp`` to the :class:`~fpgm.tracking.base.PointTracker` ABC.

    TAPNext++ cannot add query points to an already-rolling recurrent state --
    every point must be seeded on the same frame. When ``query_frame_idx > 0``,
    this class does **not** pretend to track backwards: frames strictly before
    the seed frame are reported with ``visible=False`` and ``uv`` filled with
    NaN, so a caller can never mistake "not tracked yet" for "tracked but
    occluded". If you need genuine coverage of frames before the seed, run a
    second, independent rollout seeded on an earlier frame.
    """

    def __init__(self, config: TrackingConfig, checkpoint_path: str, device: str = "cuda") -> None:
        """Args:
        config: Tracking hyperparameters (input resolution, autocast, ...).
        checkpoint_path: Path to the ``tapnextpp_512.ckpt`` weights file.
        device: Torch device string, e.g. ``"cuda"`` or ``"cuda:1"``.
        """
        self.config = config
        self.checkpoint_path = checkpoint_path
        self.device = device
        self._model = None  # built lazily so importing this module needs no torch/tapnet

    def _build(self):
        """Construct the TAPNext++ model on first use and cache it."""
        if self._model is None:
            # Heavy, GPU-framework-dependent import: kept out of module scope so
            # `import fpgm.tracking.tapnext` works in an environment with no
            # tapnet/torch installed (e.g. for unit-testing fpgm.tracking.sampling).
            from tapnet.tapnextpp.votsp2026.model import TAPNextPP

            logger.info("loading TAPNext++ checkpoint from %s", self.checkpoint_path)
            self._model = TAPNextPP.from_checkpoint(
                self.checkpoint_path,
                device=self.device,
                input_resolution=self.config.input_resolution,
            )
        return self._model

    def track(
        self,
        frames: Iterable[np.ndarray],
        query_points_xy: np.ndarray,
        query_frame_idx: int = 0,
        frame_indices: np.ndarray | None = None,
    ) -> Track2D:
        """Track ``query_points_xy`` across ``frames``, seeding on ``query_frame_idx``.

        ``frames`` is consumed lazily, one image at a time, so a long clip is
        never fully materialised in RAM -- only the current frame and the
        recurrent state are held at once.

        Args:
            frames: Iterable of ``(H, W, 3) uint8`` RGB images, in video order,
                starting at (local index 0 corresponds to) the first frame this
                call should cover.
            query_points_xy: ``(Q, 2)`` float array of ``[x, y]`` display-pixel
                query points, valid on the frame at local index
                ``query_frame_idx``.
            query_frame_idx: Local index (into ``frames``, zero-based) of the
                frame the queries are seeded on. Every earlier frame is reported
                as untracked (``visible=False``), never back-filled.
            frame_indices: Absolute video frame index for each element of
                ``frames``, in the same order. Defaults to ``0, 1, 2, ...`` (the
                local index) when not given.

        Returns:
            A :class:`~fpgm.types.Track2D` spanning every frame consumed from
            ``frames``, with ``resolution`` set from the actual frame shape.

        Raises:
            ValueError: If ``query_points_xy`` is not ``(Q, 2)``, or ``frames``
                is exhausted before reaching ``query_frame_idx``.
        """
        query_points_xy = np.asarray(query_points_xy, dtype=np.float32)
        if query_points_xy.ndim != 2 or query_points_xy.shape[1] != 2:
            raise ValueError(f"query_points_xy must be shape (Q, 2), got {query_points_xy.shape}")
        n_queries = query_points_xy.shape[0]
        if query_frame_idx < 0:
            raise ValueError(f"query_frame_idx must be >= 0, got {query_frame_idx}")

        model = self._build()

        positions: list[np.ndarray] = []
        visibilities: list[np.ndarray] = []
        abs_indices: list[int] = []
        resolution: tuple[int, int] | None = None
        state = None
        seeded = False

        for local_idx, frame_rgb in enumerate(frames):
            abs_idx = int(frame_indices[local_idx]) if frame_indices is not None else local_idx
            abs_indices.append(abs_idx)
            if resolution is None:
                resolution = (int(frame_rgb.shape[1]), int(frame_rgb.shape[0]))

            if local_idx < query_frame_idx:
                # Before the seed frame: TAPNext++ has no state yet, so there is
                # nothing to report other than "not tracked".
                positions.append(np.full((n_queries, 2), np.nan, dtype=np.float32))
                visibilities.append(np.zeros(n_queries, dtype=bool))
                continue

            frame_bgr = _rgb_to_bgr(frame_rgb)  # TAPNext++ expects OpenCV BGR layout
            if not seeded:
                pos, vis, state = model.track_frame(
                    frame_bgr,
                    query_points_xy=query_points_xy,
                    state=None,
                    autocast=self.config.autocast,
                )
                seeded = True
            else:
                pos, vis, state = model.track_frame(
                    frame_bgr, state=state, autocast=self.config.autocast
                )
            positions.append(np.asarray(pos, dtype=np.float32))
            visibilities.append(np.asarray(vis, dtype=bool))

        if not seeded:
            raise ValueError(
                f"query_frame_idx={query_frame_idx} was never reached: `frames` "
                f"yielded only {len(abs_indices)} frame(s)"
            )
        assert resolution is not None  # seeded implies at least one frame was consumed

        return Track2D(
            point_id=np.arange(n_queries, dtype=np.int32),
            frames=np.asarray(abs_indices, dtype=np.int32),
            uv=np.stack(positions, axis=0),
            visible=np.stack(visibilities, axis=0),
            resolution=resolution,
        )

    def track_video_file(
        self,
        path: str,
        query_points_xy: np.ndarray,
        query_frame_idx: int = 0,
        frame_indices: np.ndarray | None = None,
    ) -> Track2D:
        """Convenience wrapper: track query points straight from a video file.

        Frames are decoded lazily with OpenCV, one at a time, rather than
        loading the whole clip up front.

        Args:
            path: Path to an mp4 (or any OpenCV-readable video).
            query_points_xy: See :meth:`track`.
            query_frame_idx: See :meth:`track`.
            frame_indices: See :meth:`track`.

        Returns:
            A :class:`~fpgm.types.Track2D` covering the whole video.
        """
        return self.track(
            frames=self._iter_video_frames(path),
            query_points_xy=query_points_xy,
            query_frame_idx=query_frame_idx,
            frame_indices=frame_indices,
        )

    @staticmethod
    def _iter_video_frames(path: str) -> Iterator[np.ndarray]:
        """Yield RGB frames from a video file one at a time via OpenCV."""
        import cv2  # local import: keeps this helper's cost out of module load

        cap = cv2.VideoCapture(path)
        if not cap.isOpened():
            raise FileNotFoundError(f"could not open video file: {path}")
        try:
            while True:
                ok, frame_bgr = cap.read()
                if not ok:
                    break
                yield frame_bgr[..., ::-1]  # cv2 decodes BGR; PointTracker.track wants RGB in
        finally:
            cap.release()


def _rgb_to_bgr(frame_rgb: np.ndarray) -> np.ndarray:
    """Reverse the channel axis: our ABC contract is RGB in, TAPNext++ wants BGR."""
    return np.ascontiguousarray(frame_rgb[..., ::-1])