70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, List, Tuple
|
|
|
|
import numpy as np
|
|
from scipy.optimize import linear_sum_assignment
|
|
|
|
from supervision.detection.utils import box_iou_batch
|
|
|
|
if TYPE_CHECKING:
|
|
from supervision.tracker.byte_tracker.core import STrack
|
|
|
|
|
|
def indices_to_matches(
|
|
cost_matrix: np.ndarray, indices: np.ndarray, thresh: float
|
|
) -> Tuple[np.ndarray, tuple, tuple]:
|
|
matched_cost = cost_matrix[tuple(zip(*indices))]
|
|
matched_mask = matched_cost <= thresh
|
|
|
|
matches = indices[matched_mask]
|
|
unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
|
|
unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
|
|
return matches, unmatched_a, unmatched_b
|
|
|
|
|
|
def linear_assignment(
|
|
cost_matrix: np.ndarray, thresh: float
|
|
) -> Tuple[np.ndarray, Tuple[int], Tuple[int, int]]:
|
|
if cost_matrix.size == 0:
|
|
return (
|
|
np.empty((0, 2), dtype=int),
|
|
tuple(range(cost_matrix.shape[0])),
|
|
tuple(range(cost_matrix.shape[1])),
|
|
)
|
|
|
|
cost_matrix[cost_matrix > thresh] = thresh + 1e-4
|
|
row_ind, col_ind = linear_sum_assignment(cost_matrix)
|
|
indices = np.column_stack((row_ind, col_ind))
|
|
|
|
return indices_to_matches(cost_matrix, indices, thresh)
|
|
|
|
|
|
def iou_distance(atracks: List[STrack], btracks: List[STrack]) -> np.ndarray:
|
|
if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
|
|
len(btracks) > 0 and isinstance(btracks[0], np.ndarray)
|
|
):
|
|
atlbrs = atracks
|
|
btlbrs = btracks
|
|
else:
|
|
atlbrs = [track.tlbr for track in atracks]
|
|
btlbrs = [track.tlbr for track in btracks]
|
|
|
|
_ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
|
|
if _ious.size != 0:
|
|
_ious = box_iou_batch(np.asarray(atlbrs), np.asarray(btlbrs))
|
|
cost_matrix = 1 - _ious
|
|
|
|
return cost_matrix
|
|
|
|
|
|
def fuse_score(cost_matrix: np.ndarray, stracks: List[STrack]) -> np.ndarray:
|
|
if cost_matrix.size == 0:
|
|
return cost_matrix
|
|
iou_sim = 1 - cost_matrix
|
|
det_scores = np.array([strack.score for strack in stracks])
|
|
det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
|
|
fuse_sim = iou_sim * det_scores
|
|
fuse_cost = 1 - fuse_sim
|
|
return fuse_cost
|