File size: 7,022 Bytes
9b92c75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from scipy.optimize import linear_sum_assignment
def _box_to_state(box: np.ndarray) -> np.ndarray:
x0, y0, x1, y1 = box
width, height = x1 - x0, y1 - y0
cx, cy = x0 + width / 2.0, y0 + height / 2.0
scale = width * height
aspect = width / max(height, 1e-6)
return np.array([cx, cy, scale, aspect], dtype=np.float64)
def _state_to_box(state: np.ndarray) -> np.ndarray:
cx, cy, scale, aspect = state[:4]
scale = max(scale, 1e-6)
width = np.sqrt(scale * aspect)
height = scale / max(width, 1e-6)
return np.array(
[cx - width / 2.0, cy - height / 2.0, cx + width / 2.0, cy + height / 2.0],
dtype=np.float64,
)
def iou_matrix(boxes_a: np.ndarray, boxes_b: np.ndarray) -> np.ndarray:
if boxes_a.shape[0] == 0 or boxes_b.shape[0] == 0:
return np.zeros((boxes_a.shape[0], boxes_b.shape[0]), dtype=np.float64)
area_a = (boxes_a[:, 2] - boxes_a[:, 0]) * (boxes_a[:, 3] - boxes_a[:, 1])
area_b = (boxes_b[:, 2] - boxes_b[:, 0]) * (boxes_b[:, 3] - boxes_b[:, 1])
top_left = np.maximum(boxes_a[:, None, :2], boxes_b[None, :, :2])
bottom_right = np.minimum(boxes_a[:, None, 2:], boxes_b[None, :, 2:])
width_height = np.clip(bottom_right - top_left, 0, None)
intersection = width_height[..., 0] * width_height[..., 1]
union = area_a[:, None] + area_b[None, :] - intersection
return intersection / np.clip(union, 1e-9, None)
class _KalmanBoxTracker:
"""Constant-velocity Kalman filter over (cx, cy, scale, aspect) — the classic SORT state.
Aspect ratio is treated as constant (no velocity term for it), matching Bewley et al. 2016.
"""
_next_id = 1
def __init__(self, box: np.ndarray, label: int, score: float) -> None:
self.state = np.zeros(7, dtype=np.float64)
self.state[:4] = _box_to_state(box)
self.covariance = np.eye(7) * 10.0
self.covariance[4:, 4:] *= 1000.0
self._transition = np.eye(7)
for i in range(3):
self._transition[i, i + 4] = 1.0
self._observation = np.zeros((4, 7))
self._observation[:4, :4] = np.eye(4)
self._process_noise = np.eye(7) * 1.0
self._process_noise[4:, 4:] *= 0.01
self._measurement_noise = np.eye(4) * 1.0
self.id = _KalmanBoxTracker._next_id
_KalmanBoxTracker._next_id += 1
self.label = label
self.score = score
self.hits = 1
self.age = 0
self.time_since_update = 0
def predict(self) -> np.ndarray:
self.state = self._transition @ self.state
self.covariance = self._transition @ self.covariance @ self._transition.T + self._process_noise
self.age += 1
self.time_since_update += 1
state = self.state.copy()
state[2] = max(state[2], 1e-6)
return _state_to_box(state)
def update(self, box: np.ndarray, label: int, score: float) -> None:
measurement = _box_to_state(box)
innovation = measurement - self._observation @ self.state
innovation_cov = self._observation @ self.covariance @ self._observation.T + self._measurement_noise
kalman_gain = self.covariance @ self._observation.T @ np.linalg.inv(innovation_cov)
self.state = self.state + kalman_gain @ innovation
self.covariance = (np.eye(7) - kalman_gain @ self._observation) @ self.covariance
self.label = label
self.score = score
self.hits += 1
self.time_since_update = 0
def current_box(self) -> np.ndarray:
return _state_to_box(self.state)
@dataclass
class Track:
id: int
box: tuple[float, float, float, float]
label: int
score: float
hits: int
age: int
class SortTracker:
"""Minimal SORT-style tracker: Kalman motion prediction + IoU/Hungarian association.
This sits *outside* the model as a post-processing layer over independent
per-frame detections — ObjectModel-v1 itself has no temporal component.
Gives detections a persistent id and lets a track survive a few frames of
missed detection (occlusion, a confidence dip) via `max_age`.
"""
def __init__(self, iou_threshold: float = 0.3, max_age: int = 5, min_hits: int = 3) -> None:
self.iou_threshold = iou_threshold
self.max_age = max_age
self.min_hits = min_hits
self._trackers: list[_KalmanBoxTracker] = []
def update(self, boxes: np.ndarray, labels: np.ndarray, scores: np.ndarray) -> list[Track]:
"""Advance one frame. `boxes` is [N, 4] xyxy, `labels`/`scores` are [N]."""
predicted = (
np.array([tracker.predict() for tracker in self._trackers])
if self._trackers
else np.zeros((0, 4))
)
matches, unmatched_detections, _ = self._associate(predicted, boxes)
for det_idx, trk_idx in matches:
self._trackers[trk_idx].update(boxes[det_idx], int(labels[det_idx]), float(scores[det_idx]))
for det_idx in unmatched_detections:
self._trackers.append(
_KalmanBoxTracker(boxes[det_idx], int(labels[det_idx]), float(scores[det_idx]))
)
self._trackers = [t for t in self._trackers if t.time_since_update <= self.max_age]
results = []
for tracker in self._trackers:
confirmed = tracker.hits >= self.min_hits or tracker.age <= self.min_hits
if tracker.time_since_update == 0 and confirmed:
x0, y0, x1, y1 = tracker.current_box()
results.append(
Track(
id=tracker.id,
box=(x0, y0, x1, y1),
label=tracker.label,
score=tracker.score,
hits=tracker.hits,
age=tracker.age,
)
)
return results
def _associate(
self, predicted: np.ndarray, detections: np.ndarray
) -> tuple[list[tuple[int, int]], list[int], list[int]]:
if predicted.shape[0] == 0 or detections.shape[0] == 0:
return [], list(range(detections.shape[0])), list(range(predicted.shape[0]))
iou = iou_matrix(detections, predicted)
row_idx, col_idx = linear_sum_assignment(1.0 - iou)
matches: list[tuple[int, int]] = []
matched_detections: set[int] = set()
matched_trackers: set[int] = set()
for row, col in zip(row_idx, col_idx, strict=True):
if iou[row, col] >= self.iou_threshold:
matches.append((row, col))
matched_detections.add(row)
matched_trackers.add(col)
unmatched_detections = [i for i in range(detections.shape[0]) if i not in matched_detections]
unmatched_trackers = [i for i in range(predicted.shape[0]) if i not in matched_trackers]
return matches, unmatched_detections, unmatched_trackers
|