111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
"""轨道（Tracklet）与基础几何工具。
|
||
|
||
对应 C++ 参考实现 reference-c++/target_tracker/include/target_tracker/tracklet.h
|
||
（规格：docs/视觉追踪技术路线.md §3.2 / §5.2 / §7）。
|
||
"""
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Optional
|
||
|
||
from .kalman_box_filter import BoxCXCYWH, KalmanBoxFilter
|
||
|
||
|
||
@dataclass
|
||
class BoxXYXY:
|
||
"""轴对齐边界框（像素坐标，左上/右下角表示）。"""
|
||
|
||
x1: float = 0.0
|
||
y1: float = 0.0
|
||
x2: float = 0.0
|
||
y2: float = 0.0
|
||
|
||
|
||
def to_center_size(box: BoxXYXY) -> BoxCXCYWH:
|
||
# §7 TRACKLET_UPDATE：cx=(x1+x2)/2；cy=(y1+y2)/2；w=x2−x1；h=y2−y1
|
||
return BoxCXCYWH((box.x1 + box.x2) / 2.0, (box.y1 + box.y2) / 2.0,
|
||
box.x2 - box.x1, box.y2 - box.y1)
|
||
|
||
|
||
def to_xyxy(box: BoxCXCYWH) -> BoxXYXY:
|
||
# §7 TRACKLET_PREDICT 的 center_size 逆变换
|
||
return BoxXYXY(box.cx - box.w / 2.0, box.cy - box.h / 2.0,
|
||
box.cx + box.w / 2.0, box.cy + box.h / 2.0)
|
||
|
||
|
||
def compute_iou(a: BoxXYXY, b: BoxXYXY) -> float:
|
||
"""标准轴对齐 IoU（§5.2）：任一框为空、交集 ≤ 0 或并集 ≤ 0 时返回 0。"""
|
||
area_a = (a.x2 - a.x1) * (a.y2 - a.y1)
|
||
area_b = (b.x2 - b.x1) * (b.y2 - b.y1)
|
||
if area_a <= 0.0 or area_b <= 0.0:
|
||
return 0.0
|
||
|
||
inter_w = min(a.x2, b.x2) - max(a.x1, b.x1)
|
||
inter_h = min(a.y2, b.y2) - max(a.y1, b.y1)
|
||
if inter_w <= 0.0 or inter_h <= 0.0:
|
||
return 0.0
|
||
|
||
inter_area = inter_w * inter_h
|
||
union_area = area_a + area_b - inter_area
|
||
if union_area <= 0.0:
|
||
return 0.0
|
||
return inter_area / union_area
|
||
|
||
|
||
class Tracklet:
|
||
"""单条轨道：生命周期计数 + 辅助状态 + 卡尔曼滤波器（§3.2）。
|
||
|
||
构造即创建轨道（§7 步骤 5）：用首次观测初始化滤波器，age=1、hits=1、
|
||
lost_count=0，last_center / last_box 立即由观测重建（创建当帧即可参与输出）。
|
||
"""
|
||
|
||
def __init__(self, track_id: int, measurement: BoxCXCYWH, conf: float,
|
||
cls_id: int, max_lost_frames: int) -> None:
|
||
self.track_id = track_id
|
||
self.kf = KalmanBoxFilter()
|
||
self.lost_count = 0
|
||
self.max_lost = max_lost_frames
|
||
self.confidence = conf
|
||
self.cls_id = cls_id
|
||
self.age = 1
|
||
self.hits = 1
|
||
self.last_center: tuple[float, float] = (measurement.cx, measurement.cy)
|
||
self.last_box: Optional[BoxXYXY] = to_xyxy(measurement)
|
||
# 创建即源于实测检测：同步记录原始检测框（Coast 期间保留不更新）
|
||
self.last_raw_box: Optional[BoxXYXY] = self.last_box
|
||
self.kf.initiate(measurement)
|
||
|
||
@property
|
||
def is_dead(self) -> bool:
|
||
# §6.1 注记 2：lost_count > max_lost（严格大于）
|
||
return self.lost_count > self.max_lost
|
||
|
||
@property
|
||
def is_coasting(self) -> bool:
|
||
# §3.2：lost_count > 0 即处于纯预测状态
|
||
return self.lost_count > 0
|
||
|
||
def predict(self) -> None:
|
||
# §7 TRACKLET_PREDICT：预测成功时更新 last_center / last_box 且 age+1
|
||
pred = self.kf.predict()
|
||
if pred is not None:
|
||
self.last_center = (pred.cx, pred.cy)
|
||
self.last_box = to_xyxy(pred)
|
||
self.age += 1
|
||
|
||
def update(self, box: BoxXYXY, conf: float, cls_id: int) -> None:
|
||
# §7 TRACKLET_UPDATE：框 → [cx,cy,w,h] → kf.update
|
||
self.kf.update(to_center_size(box))
|
||
self.last_center = ((box.x1 + box.x2) / 2.0, (box.y1 + box.y2) / 2.0)
|
||
self.last_box = box
|
||
self.last_raw_box = box
|
||
self.confidence = conf
|
||
self.cls_id = cls_id
|
||
self.lost_count = 0
|
||
self.hits += 1
|
||
self.age += 1
|
||
|
||
def mark_miss(self) -> None:
|
||
# §7 步骤 4：未匹配计丢失
|
||
self.lost_count += 1
|
||
self.age += 1
|