forked from HQU-gxy/CVTH3PE
crosstyan
6194f083cb
- Added a new `detections` attribute to the `AffinityResult` class to store detection objects used in affinity matrix calculations. - Introduced a `tracking_detections` method to yield pairs of trackings and their corresponding detections, enhancing the functionality of the class. - Updated type hints for improved clarity and consistency in the class definition.
38 lines
941 B
Python
38 lines
941 B
Python
from dataclasses import dataclass
|
|
from typing import Sequence, Callable, Generator
|
|
|
|
from app.camera import Detection
|
|
from playground 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])
|