crosstyan
65cc646927
- Added a new `Tracking` class to encapsulate tracking data, including keypoints and velocity, with a method for predicting 3D poses based on velocity. - Refactored the pose prediction logic in `calculate_camera_affinity_matrix_jax` to utilize the new `predict` method from the `Tracking` class, enhancing clarity and modularity. - Introduced an `AffinityResult` class to manage results of affinity computations, including trackings and detections, improving the structure of the affinity calculation process. - Updated imports and type hints for better organization and consistency across the codebase.
38 lines
932 B
Python
38 lines
932 B
Python
from dataclasses import dataclass
|
|
from typing import Sequence, Callable, Generator
|
|
|
|
from app.camera import Detection
|
|
from . import Tracking
|
|
from beartype.typing import Sequence, Mapping
|
|
from jaxtyping import jaxtyped, Float, Int
|
|
from jax import Array
|
|
|
|
|
|
@dataclass
|
|
class AffinityResult:
|
|
"""
|
|
Result of affinity computation between trackings and detections.
|
|
"""
|
|
|
|
matrix: Float[Array, "T D"]
|
|
"""
|
|
Affinity matrix between trackings and detections.
|
|
"""
|
|
|
|
trackings: Sequence[Tracking]
|
|
"""
|
|
Trackings used to compute the affinity matrix.
|
|
"""
|
|
|
|
detections: Sequence[Detection]
|
|
"""
|
|
Detections used to compute the affinity matrix.
|
|
"""
|
|
|
|
indices_T: Sequence[int]
|
|
indices_D: Sequence[int]
|
|
|
|
def tracking_detections(self) -> Generator[tuple[Tracking, Detection]]:
|
|
for t, d in zip(self.indices_T, self.indices_D):
|
|
yield (self.trackings[t], self.detections[d])
|