Spaces:
Paused
Paused
File size: 7,544 Bytes
729628a f6a5256 729628a | 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 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | """Temporal tracking with Kalman filter and Hungarian association."""
from __future__ import annotations
import json
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import numpy as np
from sensorflow.schemas.unified_frame import Object3D
@dataclass
class KalmanTrack:
track_id: int
state: np.ndarray # [x, y, vx, vy]
covariance: np.ndarray
missed_frames: int = 0
history: List[Dict] = field(default_factory=list)
class_name: str = "unknown"
confidence: float = 0.0
def predict(self, dt: float = 0.1):
F = np.array([[1, 0, dt, 0], [0, 1, 0, dt], [0, 0, 1, 0], [0, 0, 0, 1]])
Q = np.eye(4) * 0.1
self.state = F @ self.state
self.covariance = F @ self.covariance @ F.T + Q
def update(self, measurement: np.ndarray):
H = np.array([[1, 0, 0, 0], [0, 1, 0, 0]])
R = np.eye(2) * 0.5
y = measurement - H @ self.state
S = H @ self.covariance @ H.T + R
K = self.covariance @ H.T @ np.linalg.inv(S)
self.state = self.state + K @ y
self.covariance = (np.eye(4) - K @ H) @ self.covariance
self.missed_frames = 0
class TemporalTracker:
"""Kalman + Hungarian tracker for 3D bounding box sequences."""
def __init__(
self,
max_age: int = 30,
distance_gate: float = 3.0,
velocity_penalty: float = 2.0,
):
self.max_age = max_age
self.distance_gate = distance_gate
self.velocity_penalty = velocity_penalty
self.tracks: Dict[int, KalmanTrack] = {}
self.next_id = 1
self.id_swap_log: List[Dict] = []
def update(self, frame_id: str, proposals: List[Object3D]) -> List[Object3D]:
for track in self.tracks.values():
track.predict()
track.missed_frames += 1
if proposals and self.tracks:
assignments = self._associate(proposals)
matched_proposals = set()
for track_id, prop_idx in assignments:
prop = proposals[prop_idx]
track = self.tracks[track_id]
measurement = np.array([prop.bbox_3d[0], prop.bbox_3d[1]])
track.update(measurement)
track.class_name = prop.class_name
track.confidence = prop.confidence
track.history.append({
"frame_id": frame_id,
"bbox_3d": prop.bbox_3d,
"confidence": prop.confidence,
})
matched_proposals.add(prop_idx)
for i, prop in enumerate(proposals):
if i not in matched_proposals:
self._spawn_track(frame_id, prop)
elif proposals:
for prop in proposals:
self._spawn_track(frame_id, prop)
dead = [tid for tid, t in self.tracks.items() if t.missed_frames > self.max_age]
for tid in dead:
del self.tracks[tid]
result = []
for track in self.tracks.values():
if track.history and track.history[-1]["frame_id"] == frame_id:
prop = Object3D(
bbox_3d=track.history[-1]["bbox_3d"],
class_name=track.class_name,
confidence=track.confidence,
track_id=track.track_id,
)
result.append(prop)
return result
def _associate(self, proposals: List[Object3D]) -> List[Tuple[int, int]]:
try:
from scipy.optimize import linear_sum_assignment
except ImportError:
return self._greedy_associate(proposals)
track_ids = list(self.tracks.keys())
cost = np.full((len(track_ids), len(proposals)), 1e6)
for ti, tid in enumerate(track_ids):
track = self.tracks[tid]
for pi, prop in enumerate(proposals):
dist = math.sqrt(
(track.state[0] - prop.bbox_3d[0]) ** 2 +
(track.state[1] - prop.bbox_3d[1]) ** 2
)
# Yaw penalty: compare against the track's last *observed* yaw
# (the Kalman state is [x, y, vx, vy] and carries no yaw).
# The difference is wrapped to [-pi, pi] so 359 deg vs 1 deg
# is a 2-degree disagreement, not a 358-degree one.
if track.history:
last_yaw = track.history[-1]["bbox_3d"][6]
raw = last_yaw - prop.bbox_3d[6]
yaw_diff = abs(math.atan2(math.sin(raw), math.cos(raw)))
else:
yaw_diff = 0.0
vel_penalty = self.velocity_penalty * yaw_diff
cost[ti, pi] = dist + vel_penalty
row_ind, col_ind = linear_sum_assignment(cost)
assignments = []
for r, c in zip(row_ind, col_ind):
if cost[r, c] < self.distance_gate:
assignments.append((track_ids[r], c))
return assignments
def _greedy_associate(self, proposals: List[Object3D]) -> List[Tuple[int, int]]:
assignments = []
used = set()
for tid, track in self.tracks.items():
best_pi, best_dist = None, float("inf")
for pi, prop in enumerate(proposals):
if pi in used:
continue
dist = math.sqrt(
(track.state[0] - prop.bbox_3d[0]) ** 2 +
(track.state[1] - prop.bbox_3d[1]) ** 2
)
if dist < best_dist and dist < self.distance_gate:
best_dist = dist
best_pi = pi
if best_pi is not None:
assignments.append((tid, best_pi))
used.add(best_pi)
return assignments
def _spawn_track(self, frame_id: str, prop: Object3D):
state = np.array([prop.bbox_3d[0], prop.bbox_3d[1], 0.0, 0.0])
track = KalmanTrack(
track_id=self.next_id,
state=state,
covariance=np.eye(4),
class_name=prop.class_name,
confidence=prop.confidence,
)
track.history.append({
"frame_id": frame_id,
"bbox_3d": prop.bbox_3d,
"confidence": prop.confidence,
})
self.tracks[self.next_id] = track
self.next_id += 1
def run_sequence(
self,
proposals_by_frame: Dict[str, List[Object3D]],
output_path: Path,
) -> List[Dict]:
for frame_id in sorted(proposals_by_frame.keys()):
self.update(frame_id, proposals_by_frame[frame_id])
tracks_output = []
for tid, track in self.tracks.items():
if track.history:
tracks_output.append({
"track_id": tid,
"class_name": track.class_name,
"frames": track.history,
})
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w") as f:
json.dump(tracks_output, f, indent=2)
return tracks_output
@staticmethod
def load_proposals(proposals_dir: Path) -> Dict[str, List[Object3D]]:
result = {}
for p in sorted(proposals_dir.glob("*.json")):
frame_id = p.stem
with open(p) as f:
data = json.load(f)
result[frame_id] = [Object3D.model_validate(d) for d in data]
return result
|