first commit

2026-09-02 15:13:13 +08:00
commit cb00322a50
107 changed files with 10918 additions and 0 deletions
@@ -0,0 +1,269 @@
+"""多目标追踪过滤器（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