src/fpgm/tracking/base.py

"""Abstract interface for point trackers."""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Iterable

import numpy as np

from fpgm.types import Track2D


class PointTracker(ABC):
    """Tracks a fixed set of query points across a video.

    The fixed-set contract is deliberate: TAPNext++ cannot add new queries to an
    existing recurrent state, so all query points must be seeded on the same frame
    (or a separate rollout run per seed frame).
    """

    @abstractmethod
    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`` (Q, 2) in display pixels across ``frames``.

        ``frames`` yields ``(H, W, 3) uint8`` images. Implementations are
        responsible for whatever colour order their backend needs; callers always
        pass RGB.
        """