CVTH3PE/app/tracking/__init__.py

from dataclasses import dataclass
from datetime import datetime
from typing import (
    Any,
    Callable,
    Generator,
    Optional,
    Sequence,
    TypeAlias,
    TypedDict,
    TypeVar,
    cast,
    overload,
    Protocol,
)
from datetime import timedelta
import jax.numpy as jnp
from beartype import beartype
from beartype.typing import Mapping, Sequence
from jax import Array
from jaxtyping import Array, Float, Int, jaxtyped
from pyrsistent import PVector, v
from itertools import chain
from app.camera import Detection


class TrackingPrediction(TypedDict):
    velocity: Float[Array, "J 3"]
    keypoints: Float[Array, "J 3"]


class GenericVelocityFilter(Protocol):
    """
    a filter interface for tracking velocity estimation
    """

    def predict(self, timestamp: datetime) -> TrackingPrediction:
        """
        predict the velocity and the keypoints location

        Args:
            timestamp: timestamp of the prediction

        Returns:
            velocity: velocity of the tracking
            keypoints: keypoints of the tracking
        """
        ...  # pylint: disable=unnecessary-ellipsis

    def update(self, keypoints: Float[Array, "J 3"], timestamp: datetime) -> None:
        """
        update the filter state with new measurements

        Args:
            keypoints: new measurements
            timestamp: timestamp of the update
        """
        ...  # pylint: disable=unnecessary-ellipsis

    def get(self) -> TrackingPrediction:
        """
        get the current state of the filter state

        Returns:
            velocity: velocity of the tracking
            keypoints: keypoints of the tracking
        """
        ...  # pylint: disable=unnecessary-ellipsis

    def reset(self, keypoints: Float[Array, "J 3"], timestamp: datetime) -> None:
        """
        reset the filter state with new keypoints

        Args:
            keypoints: new keypoints
            timestamp: timestamp of the reset
        """
        ...  # pylint: disable=unnecessary-ellipsis


class LastDifferenceVelocityFilter(GenericVelocityFilter):
    """
    a naive velocity filter that uses the last difference of keypoints
    """

    _last_timestamp: datetime
    _last_keypoints: Float[Array, "J 3"]
    _last_velocity: Optional[Float[Array, "J 3"]] = None

    def __init__(self, keypoints: Float[Array, "J 3"], timestamp: datetime):
        self._last_keypoints = keypoints
        self._last_timestamp = timestamp

    def predict(self, timestamp: datetime) -> TrackingPrediction:
        delta_t_s = (timestamp - self._last_timestamp).total_seconds()
        if self._last_velocity is None:
            return TrackingPrediction(
                velocity=jnp.zeros_like(self._last_keypoints),
                keypoints=self._last_keypoints,
            )
        else:
            return TrackingPrediction(
                velocity=self._last_velocity,
                keypoints=self._last_keypoints + self._last_velocity * delta_t_s,
            )

    def update(self, keypoints: Float[Array, "J 3"], timestamp: datetime) -> None:
        delta_t_s = (timestamp - self._last_timestamp).total_seconds()
        self._last_velocity = (keypoints - self._last_keypoints) / delta_t_s
        self._last_keypoints = keypoints
        self._last_timestamp = timestamp

    def get(self) -> TrackingPrediction:
        if self._last_velocity is None:
            return TrackingPrediction(
                velocity=jnp.zeros_like(self._last_keypoints),
                keypoints=self._last_keypoints,
            )
        else:
            return TrackingPrediction(
                velocity=self._last_velocity,
                keypoints=self._last_keypoints,
            )

    def reset(self, keypoints: Float[Array, "J 3"], timestamp: datetime) -> None:
        self._last_keypoints = keypoints
        self._last_timestamp = timestamp
        self._last_velocity = None


class LeastMeanSquareVelocityFilter(GenericVelocityFilter):
    _get_historical_detections: Callable[[], Sequence[Detection]]
    """
    get the current historical detections, assuming the detections are sorted by
    timestamp incrementally (i.e. index 0 is the oldest detection, index -1 is
    the newest detection)
    """
    _velocity: Optional[Float[Array, "J 3"]] = None

    def __init__(self, get_historical_detections: Callable[[], Sequence[Detection]]):
        self._get_historical_detections = get_historical_detections

    def predict(self, timestamp: datetime) -> TrackingPrediction:
        raise NotImplementedError

    @jaxtyped(typechecker=beartype)
    def _update(
        self,
        keypoints: Float[Array, "N J 3"],
        timestamps: Float[Array, "N"],
    ) -> None:
        """
        update measurements with least mean square method
        """
        raise NotImplementedError

    def update(self, keypoints: Float[Array, "J 3"], timestamp: datetime) -> None:
        historical_detections = self._get_historical_detections()
        t_0 = min(d.timestamp for d in historical_detections)
        detections = jnp.array(
            chain((d.keypoints for d in historical_detections), (keypoints,))
        )
        # timestamps relative to t_0 (the oldest detection timestamp)
        timestamps = jnp.array(
            chain(
                ((d.timestamp - t_0).total_seconds() for d in historical_detections),
                ((timestamp - t_0).total_seconds(),),
            )
        )
        raise NotImplementedError

    def get(self) -> TrackingPrediction:
        raise NotImplementedError

    def reset(self, keypoints: Float[Array, "J 3"], timestamp: datetime) -> None:
        raise NotImplementedError


@jaxtyped(typechecker=beartype)
@dataclass(frozen=True)
class Tracking:
    id: int
    """
    The tracking id
    """
    keypoints: Float[Array, "J 3"]
    """
    The 3D keypoints of the tracking

    Used for calculate affinity 3D
    """
    last_active_timestamp: datetime
    """
    The last active timestamp of the tracking
    """

    historical_detections: PVector[Detection]
    """
    Historical detections of the tracking.

    Used for 3D re-triangulation
    """

    velocity_filter: GenericVelocityFilter
    """
    The velocity filter of the tracking
    """

    def __repr__(self) -> str:
        return f"Tracking({self.id}, {self.last_active_timestamp})"

    @overload
    def predict(self, time: float) -> Float[Array, "J 3"]:
        """
        predict the keypoints at a given time

        Args:
            time: the time in seconds to predict the keypoints

        Returns:
            the predicted keypoints
        """
        ...  # pylint: disable=unnecessary-ellipsis

    @overload
    def predict(self, time: timedelta) -> Float[Array, "J 3"]:
        """
        predict the keypoints at a given time

        Args:
            time: the time delta to predict the keypoints
        """
        ...  # pylint: disable=unnecessary-ellipsis

    @overload
    def predict(self, time: datetime) -> Float[Array, "J 3"]:
        """
        predict the keypoints at a given time

        Args:
            time: the timestamp to predict the keypoints
        """
        ...  # pylint: disable=unnecessary-ellipsis

    def predict(
        self,
        time: float | timedelta | datetime,
    ) -> Float[Array, "J 3"]:
        if isinstance(time, timedelta):
            timestamp = self.last_active_timestamp + time
        elif isinstance(time, datetime):
            timestamp = time
        else:
            timestamp = self.last_active_timestamp + timedelta(seconds=time)
        # pylint: disable-next=unsubscriptable-object
        return self.velocity_filter.predict(timestamp)["keypoints"]

    @property
    def velocity(self) -> Float[Array, "J 3"]:
        """
        The velocity of the tracking for each keypoint
        """
        # pylint: disable-next=unsubscriptable-object
        return self.velocity_filter.get()["velocity"]


@jaxtyped(typechecker=beartype)
@dataclass
class AffinityResult:
    """
    Result of affinity computation between trackings and detections.
    """

    matrix: Float[Array, "T D"]
    trackings: Sequence[Tracking]
    detections: Sequence[Detection]
    indices_T: Int[Array, "T"]  # pylint: disable=invalid-name
    indices_D: Int[Array, "D"]  # pylint: disable=invalid-name

    def tracking_detections(
        self,
    ) -> Generator[tuple[float, Tracking, Detection], None, None]:
        """
        iterate over the best matching trackings and detections
        """
        for t, d in zip(self.indices_T, self.indices_D):
            yield (
                self.matrix[t, d].item(),
                self.trackings[t],
                self.detections[d],
            )