crosstyan
da4c51d04f
- Added a new `AffinityResult` class to encapsulate the results of affinity computations, including the affinity matrix, trackings, and their respective indices. - Introduced a vectorized implementation of `calculate_camera_affinity_matrix_jax` to enhance performance by leveraging JAX's capabilities, replacing the previous double-for-loop approach. - Updated tests in `test_affinity.py` to include parameterized benchmarks for comparing the performance of the new vectorized method against the naive implementation, ensuring accuracy and efficiency.
35 lines
776 B
Python
35 lines
776 B
Python
from dataclasses import dataclass
|
|
from typing import Sequence
|
|
|
|
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.
|
|
"""
|
|
|
|
indices_T: Sequence[int]
|
|
"""
|
|
Indices of the trackings that were used to compute the affinity matrix.
|
|
"""
|
|
|
|
indices_D: Sequence[int]
|
|
"""
|
|
Indices of the detections that were used to compute the affinity matrix.
|
|
"""
|