| 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 |
|
|
