107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
"""单摄像头目标跟踪（轻量 IoU 关联）
|
||
|
||
设计：参考 Frigate/Norfair 的追踪思路，但用最小依赖实现一个 IoU 关联器，
|
||
避免强制引入 norfair。后续可平滑替换为 norfair 或 DeepSORT/ByteTrack 的特征关联。
|
||
|
||
输出：为每个检测框分配 track_id，并维护其在场状态/累计帧数/最近一帧框。
|
||
"""
|
||
import logging
|
||
|
||
logger = logging.getLogger("analysis.tracker")
|
||
|
||
try:
|
||
import numpy as np # reserved for future vectorized IoU; not required for operation
|
||
_ = np
|
||
except Exception:
|
||
np = None
|
||
|
||
|
||
def _iou(a, b):
|
||
ax1, ay1, ax2, ay2 = a
|
||
bx1, by1, bx2, by2 = b
|
||
ix1 = max(ax1, bx1); iy1 = max(ay1, by1)
|
||
ix2 = min(ax2, bx2); iy2 = min(ay2, by2)
|
||
iw = max(0, ix2 - ix1); ih = max(0, iy2 - iy1)
|
||
inter = iw * ih
|
||
a_area = max(0, ax2 - ax1) * max(0, ay2 - ay1)
|
||
b_area = max(0, bx2 - bx1) * max(0, by2 - by1)
|
||
union = a_area + b_area - inter
|
||
if union <= 0:
|
||
return 0.0
|
||
return float(inter) / float(union)
|
||
|
||
|
||
class Track(object):
|
||
__slots__ = ("track_id", "label", "box", "score", "missed", "hits", "born")
|
||
|
||
def __init__(self, track_id, label, box, score, born):
|
||
self.track_id = track_id
|
||
self.label = label
|
||
self.box = box
|
||
self.score = score
|
||
self.missed = 0
|
||
self.hits = 1
|
||
self.born = born
|
||
|
||
|
||
class IoUTracker(object):
|
||
"""按类别维护轨迹，IoU 匹配；max_missed 后判定目标消失。"""
|
||
|
||
def __init__(self, iou_threshold=0.3, max_missed=8):
|
||
self.iou_threshold = iou_threshold
|
||
self.max_missed = max_missed
|
||
self._tracks = {} # track_id -> Track
|
||
self._next_id = 1
|
||
|
||
def update(self, detections, frame_index):
|
||
"""detections: list[dict(box, label, score)]
|
||
返回 list[dict(track_id, label, box, score)] 当前帧仍在场的轨迹"""
|
||
active = {}
|
||
new_tracks = []
|
||
# 贪心匹配
|
||
for det in detections:
|
||
best_id = None
|
||
best_iou = self.iou_threshold
|
||
for tid, tr in self._tracks.items():
|
||
if tr.label != det["label"]:
|
||
continue
|
||
v = _iou(tr.box, det["box"])
|
||
if v > best_iou:
|
||
best_iou = v
|
||
best_id = tid
|
||
if best_id is not None:
|
||
tr = self._tracks[best_id]
|
||
tr.box = det["box"]
|
||
tr.score = det["score"]
|
||
tr.missed = 0
|
||
tr.hits += 1
|
||
active[best_id] = tr
|
||
else:
|
||
tid = self._next_id
|
||
self._next_id += 1
|
||
tr = Track(tid, det["label"], det["box"], det["score"], frame_index)
|
||
self._tracks[tid] = tr
|
||
active[tid] = tr
|
||
new_tracks.append(tid)
|
||
|
||
# 未匹配的轨迹累计 missed
|
||
ended = []
|
||
for tid, tr in self._tracks.items():
|
||
if tid in active:
|
||
continue
|
||
tr.missed += 1
|
||
if tr.missed >= self.max_missed:
|
||
ended.append(tid)
|
||
for tid in ended:
|
||
del self._tracks[tid]
|
||
|
||
return [{"track_id": t.track_id, "label": t.label, "box": t.box, "score": t.score}
|
||
for t in active.values()], ended, new_tracks, frame_index
|
||
|
||
def all_active(self):
|
||
return list(self._tracks.values())
|
||
|
||
def reset(self):
|
||
self._tracks.clear()
|
||
self._next_id = 1
|