270 lines
10 KiB
Python
270 lines
10 KiB
Python
"""多目标追踪过滤器（TargetTracker）。
|
||
|
||
对应 C++ 参考实现 reference-c++/target_tracker/src/target_tracker.cpp
|
||
（规格：docs/视觉追踪技术路线.md §5 关联 / §6 生命周期与锁定 / §7 六步主流程）。
|
||
|
||
线程模型（§9.5）：update() 由单线程（推理主循环）调用；仅对 last_result
|
||
缓存加锁，last_result() 返回受保护拷贝。
|
||
"""
|
||
|
||
import math
|
||
import threading
|
||
from dataclasses import dataclass, field
|
||
from typing import Optional, Sequence
|
||
|
||
from .kalman_box_filter import BoxCXCYWH
|
||
from .tracklet import BoxXYXY, Tracklet, compute_iou, to_center_size
|
||
|
||
|
||
@dataclass
|
||
class Detection:
|
||
"""单个检测框（§2.2 输入元素：[x1,y1,x2,y2,conf,cls_id]）。"""
|
||
|
||
x1: float = 0.0
|
||
y1: float = 0.0
|
||
x2: float = 0.0
|
||
y2: float = 0.0
|
||
conf: float = 0.0
|
||
cls_id: int = 0
|
||
|
||
|
||
@dataclass
|
||
class FrameShape:
|
||
"""画面尺寸（height, width）：仅 Coast 兜底选举时用于取画面中心。"""
|
||
|
||
height: int = 0
|
||
width: int = 0
|
||
|
||
|
||
@dataclass
|
||
class TrackResult:
|
||
"""每帧追踪输出（§2.3 九字段）。"""
|
||
|
||
tracked: bool = False
|
||
primary_id: Optional[int] = None
|
||
center: tuple[float, float] = (0.0, 0.0)
|
||
box: Optional[BoxXYXY] = None
|
||
confidence: float = 0.0
|
||
is_predicted: bool = False
|
||
lost_frames: int = 0
|
||
n_active: int = 0
|
||
raw: Optional[BoxXYXY] = None
|
||
|
||
|
||
@dataclass
|
||
class TrackerConfig:
|
||
"""追踪器配置参数（§8.1 默认值取自规格）。"""
|
||
|
||
max_lost_frames: int = 8
|
||
max_association_dist: float = 200.0
|
||
min_hits: int = 3
|
||
dist_weight: float = 0.5
|
||
iou_weight: float = 0.5
|
||
|
||
|
||
@dataclass
|
||
class _AssociationResult:
|
||
matches: list[tuple[int, int]] = field(default_factory=list)
|
||
unmatched_track_ids: list[int] = field(default_factory=list)
|
||
unmatched_det_indices: list[int] = field(default_factory=list)
|
||
|
||
|
||
class TargetTracker:
|
||
def __init__(self, config: Optional[TrackerConfig] = None) -> None:
|
||
self._config = config if config is not None else TrackerConfig()
|
||
self._next_id = 0
|
||
self._tracks: dict[int, Tracklet] = {}
|
||
self._frame_count = 0
|
||
self._locked_primary_id: Optional[int] = None
|
||
self._last_result = TrackResult()
|
||
self._last_result_mutex = threading.Lock()
|
||
|
||
@property
|
||
def config(self) -> TrackerConfig:
|
||
return self._config
|
||
|
||
def get_config(self) -> TrackerConfig:
|
||
return self._config
|
||
|
||
def update(self, detections: Sequence[Detection],
|
||
frame_shape: Optional[FrameShape] = None) -> TrackResult:
|
||
# §7 六步主流程
|
||
self._frame_count += 1
|
||
|
||
# 入口防御：过滤坐标或置信度含 NaN/Inf 的检测（非有限值一旦进入
|
||
# 卡尔曼状态将永久污染轨道并传导到主目标输出）
|
||
valid_detections = [
|
||
det for det in detections
|
||
if all(math.isfinite(v) for v in (det.x1, det.y1, det.x2, det.y2, det.conf))
|
||
]
|
||
|
||
# 步骤 1：所有存活轨道先行预测（推进一帧）
|
||
for track in self._tracks.values():
|
||
track.predict()
|
||
|
||
# 步骤 2：贪婪关联（无检测时全部轨道进入 Coast 计数）
|
||
if valid_detections:
|
||
assoc = self._associate_detections(valid_detections)
|
||
else:
|
||
assoc = _AssociationResult(unmatched_track_ids=list(self._tracks.keys()))
|
||
|
||
# 步骤 3：用匹配到的检测更新轨道
|
||
for track_id, det_idx in assoc.matches:
|
||
det = valid_detections[det_idx]
|
||
raw_box = BoxXYXY(det.x1, det.y1, det.x2, det.y2)
|
||
self._tracks[track_id].update(raw_box, det.conf, det.cls_id)
|
||
|
||
# 步骤 4：未匹配轨道计丢失，死亡者回收（锁定唯一解除途径之一）
|
||
for track_id in assoc.unmatched_track_ids:
|
||
track = self._tracks.get(track_id)
|
||
if track is None:
|
||
continue
|
||
track.mark_miss()
|
||
if track.is_dead:
|
||
if self._locked_primary_id == track_id:
|
||
self._locked_primary_id = None
|
||
del self._tracks[track_id]
|
||
|
||
# 步骤 5：未匹配检测 → 创建新轨道（next_id 从 0 递增永不复用）
|
||
for det_idx in assoc.unmatched_det_indices:
|
||
det = valid_detections[det_idx]
|
||
measurement = to_center_size(BoxXYXY(det.x1, det.y1, det.x2, det.y2))
|
||
self._tracks[self._next_id] = Tracklet(
|
||
self._next_id, measurement, det.conf, det.cls_id,
|
||
self._config.max_lost_frames)
|
||
self._next_id += 1
|
||
|
||
# 步骤 6：选取/延续主目标，组装输出并缓存
|
||
result = self._select_primary(frame_shape)
|
||
with self._last_result_mutex:
|
||
self._last_result = result
|
||
return result
|
||
|
||
def last_result(self) -> TrackResult:
|
||
with self._last_result_mutex:
|
||
return self._last_result
|
||
|
||
def reset(self) -> None:
|
||
# §2.1：清空所有轨道、ID 计数器、锁定状态与缓存结果（next_id 归零）
|
||
self._tracks.clear()
|
||
self._next_id = 0
|
||
self._frame_count = 0
|
||
self._locked_primary_id = None
|
||
with self._last_result_mutex:
|
||
self._last_result = TrackResult()
|
||
|
||
def active_count(self) -> int:
|
||
return len(self._tracks)
|
||
|
||
def active_tracks(self) -> dict[int, Tracklet]:
|
||
return dict(self._tracks)
|
||
|
||
def _associate_detections(self, detections: Sequence[Detection]) -> _AssociationResult:
|
||
result = _AssociationResult()
|
||
|
||
# §5.3：按轨道 confidence 降序（平局按 ID 升序保证确定性）。
|
||
# 非有限置信度按 −∞ 处理（排最后），避免 NaN 破坏全序。
|
||
def sort_key(track_id: int) -> tuple[float, int]:
|
||
raw = self._tracks[track_id].confidence
|
||
conf = raw if math.isfinite(raw) else -math.inf
|
||
return (-conf, track_id)
|
||
|
||
sorted_ids = sorted(self._tracks.keys(), key=sort_key)
|
||
det_assigned = [False] * len(detections)
|
||
|
||
for track_id in sorted_ids:
|
||
track = self._tracks[track_id]
|
||
# §5.3：最优代价初始 1.0，仅 cost 严格小于才匹配（cost ≥ 1.0 判为不可匹配）
|
||
best_cost = 1.0
|
||
best_det = -1
|
||
for d, det in enumerate(detections):
|
||
if det_assigned[d]:
|
||
continue
|
||
det_cx = (det.x1 + det.x2) / 2.0
|
||
det_cy = (det.y1 + det.y2) / 2.0
|
||
dist = math.hypot(track.last_center[0] - det_cx,
|
||
track.last_center[1] - det_cy)
|
||
norm_dist = dist / max(self._config.max_association_dist, 1e-3)
|
||
iou = (compute_iou(track.last_box,
|
||
BoxXYXY(det.x1, det.y1, det.x2, det.y2))
|
||
if track.last_box is not None else 0.0)
|
||
cost = (self._config.dist_weight * norm_dist +
|
||
self._config.iou_weight * (1.0 - iou))
|
||
if cost < best_cost:
|
||
best_cost = cost
|
||
best_det = d
|
||
if best_det >= 0:
|
||
result.matches.append((track_id, best_det))
|
||
det_assigned[best_det] = True
|
||
else:
|
||
result.unmatched_track_ids.append(track_id)
|
||
|
||
result.unmatched_det_indices = [
|
||
d for d, assigned in enumerate(det_assigned) if not assigned
|
||
]
|
||
return result
|
||
|
||
def _select_primary(self, frame_shape: Optional[FrameShape]) -> TrackResult:
|
||
result = TrackResult(n_active=len(self._tracks))
|
||
|
||
# §6.3：轨道表为空 → 清除锁定，返回 tracked=false
|
||
if not self._tracks:
|
||
self._locked_primary_id = None
|
||
return result
|
||
|
||
# §6.3 ① 锁定延续：锁定轨道存活则直接输出（Coast 中照样输出，
|
||
# 不被高置信度目标抢占）
|
||
chosen: Optional[Tracklet] = None
|
||
is_predicted = False
|
||
if self._locked_primary_id is not None:
|
||
track = self._tracks.get(self._locked_primary_id)
|
||
if track is not None:
|
||
chosen = track
|
||
is_predicted = track.is_coasting
|
||
|
||
# §6.3 ② 重新选举（无锁定 / 锁定目标已死亡被回收）
|
||
if chosen is None:
|
||
measured = [t for t in self._tracks.values() if not t.is_coasting]
|
||
coasting = [t for t in self._tracks.values() if t.is_coasting]
|
||
|
||
if measured:
|
||
# confirmed ← measured 中 hits ≥ min_hits 者；为空则退回 measured
|
||
# （min_hits 只影响候选池优先级，不阻止输出）
|
||
pool = [t for t in measured if t.hits >= self._config.min_hits]
|
||
if not pool:
|
||
pool = measured
|
||
chosen = pool[0]
|
||
for t in pool:
|
||
if t.confidence > chosen.confidence:
|
||
chosen = t
|
||
is_predicted = False
|
||
elif coasting:
|
||
# 无实测轨道：选离画面中心最近者（缺省按 640×480 中心 (320,240)）
|
||
center_x = frame_shape.width / 2.0 if frame_shape and frame_shape.width > 0 else 320.0
|
||
center_y = frame_shape.height / 2.0 if frame_shape and frame_shape.height > 0 else 240.0
|
||
best_dist_sq = math.inf
|
||
for t in coasting:
|
||
dx = t.last_center[0] - center_x
|
||
dy = t.last_center[1] - center_y
|
||
dist_sq = dx * dx + dy * dy
|
||
if dist_sq < best_dist_sq:
|
||
best_dist_sq = dist_sq
|
||
chosen = t
|
||
is_predicted = True
|
||
else:
|
||
# 理论上不可达（轨道非空则 measured/coasting 必居其一）；防御清锁
|
||
self._locked_primary_id = None
|
||
return result
|
||
self._locked_primary_id = chosen.track_id # 选举后锁定
|
||
|
||
# 组装输出（§2.3 九字段）；raw 仅实测时有值，Coast 时为空
|
||
result.tracked = True
|
||
result.primary_id = chosen.track_id
|
||
result.center = chosen.last_center
|
||
result.box = chosen.last_box
|
||
result.confidence = chosen.confidence
|
||
result.is_predicted = is_predicted
|
||
result.lost_frames = chosen.lost_count
|
||
result.raw = None if is_predicted else chosen.last_raw_box
|
||
return result
|