189 lines
7.1 KiB
Python
189 lines
7.1 KiB
Python
"""低帧率视频目标跟踪。
|
|
|
|
采用同类别全局一对一匹配，并用时间而不是丢帧数保留轨迹。这样在 CPU
|
|
推理只有约 0.3 FPS 时，短暂缺检不会不断生成新的 track_id。
|
|
"""
|
|
import math
|
|
import time
|
|
|
|
|
|
def _iou(a, b):
|
|
ax1, ay1, ax2, ay2 = a
|
|
bx1, by1, bx2, by2 = b
|
|
ix1, iy1 = max(ax1, bx1), max(ay1, by1)
|
|
ix2, iy2 = min(ax2, bx2), min(ay2, by2)
|
|
inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
|
|
area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
|
|
area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
|
|
union = area_a + area_b - inter
|
|
return float(inter / union) if union > 0 else 0.0
|
|
|
|
|
|
def _center(box):
|
|
return ((box[0] + box[2]) * 0.5, (box[1] + box[3]) * 0.5)
|
|
|
|
|
|
def _center_distance(a, b):
|
|
ac, bc = _center(a), _center(b)
|
|
distance = math.hypot(ac[0] - bc[0], ac[1] - bc[1])
|
|
diag = max(1.0, math.hypot(max(a[2] - a[0], b[2] - b[0]),
|
|
max(a[3] - a[1], b[3] - b[1])))
|
|
return distance / diag
|
|
|
|
|
|
class Track(object):
|
|
__slots__ = (
|
|
"track_id", "label", "box", "score", "missed", "hits", "born",
|
|
"last_seen", "last_update", "velocity", "confirmed", "observed",
|
|
"algorithm_id", "algorithm_name",
|
|
)
|
|
|
|
def __init__(self, track_id, detection, born, timestamp):
|
|
self.track_id = track_id
|
|
self.label = detection.get("label", "unknown")
|
|
self.box = list(detection.get("box") or [0, 0, 0, 0])
|
|
self.score = float(detection.get("score", 0) or 0)
|
|
self.missed = 0
|
|
self.hits = 1
|
|
self.born = born
|
|
self.last_seen = timestamp
|
|
self.last_update = timestamp
|
|
self.velocity = [0.0, 0.0, 0.0, 0.0]
|
|
self.confirmed = self.score >= 0.6
|
|
self.observed = True
|
|
self.algorithm_id = detection.get("algorithm_id")
|
|
self.algorithm_name = detection.get("algorithm_name", "")
|
|
|
|
def predicted_box(self, timestamp):
|
|
dt = max(0.0, min(5.0, timestamp - self.last_update))
|
|
return [self.box[i] + self.velocity[i] * dt for i in range(4)]
|
|
|
|
def as_dict(self):
|
|
return {
|
|
"track_id": self.track_id,
|
|
"label": self.label,
|
|
"box": list(self.box),
|
|
"score": self.score,
|
|
"hits": self.hits,
|
|
"confirmed": self.confirmed,
|
|
"observed": self.observed,
|
|
"last_seen": self.last_seen,
|
|
"algorithm_id": self.algorithm_id,
|
|
"algorithm_name": self.algorithm_name,
|
|
"velocity": list(self.velocity),
|
|
}
|
|
|
|
|
|
class IoUTracker(object):
|
|
"""低 FPS 下的全局一对一框关联器。"""
|
|
|
|
def __init__(self, iou_threshold=0.2, max_missed=8):
|
|
self.iou_threshold = float(iou_threshold)
|
|
self.max_missed = max_missed # 兼容旧构造参数；实际淘汰按时间
|
|
self._tracks = {}
|
|
self._next_id = 1
|
|
self._last_update_ts = None
|
|
self._recent_period = 1.0
|
|
|
|
@property
|
|
def retention_sec(self):
|
|
return max(2.0, min(15.0, 3.0 * self._recent_period))
|
|
|
|
def update(self, detections, frame_index, timestamp=None, max_age_sec=None):
|
|
timestamp = float(timestamp if timestamp is not None else time.time())
|
|
if self._last_update_ts is not None:
|
|
period = timestamp - self._last_update_ts
|
|
if 0.05 <= period <= 60:
|
|
self._recent_period = self._recent_period * 0.7 + period * 0.3
|
|
self._last_update_ts = timestamp
|
|
max_age = float(max_age_sec if max_age_sec is not None else self.retention_sec)
|
|
|
|
for track in self._tracks.values():
|
|
track.observed = False
|
|
|
|
candidates = []
|
|
for det_index, det in enumerate(detections or []):
|
|
box = det.get("box") or []
|
|
if len(box) != 4:
|
|
continue
|
|
for tid, track in self._tracks.items():
|
|
if track.label != det.get("label"):
|
|
continue
|
|
predicted = track.predicted_box(timestamp)
|
|
overlap = _iou(predicted, box)
|
|
distance = _center_distance(predicted, box)
|
|
if overlap < self.iou_threshold and distance > 1.5:
|
|
continue
|
|
affinity = overlap * 0.72 + max(0.0, 1.0 - distance / 1.5) * 0.28
|
|
candidates.append((affinity, overlap, -distance, tid, det_index))
|
|
candidates.sort(reverse=True)
|
|
|
|
used_tracks, used_detections = set(), set()
|
|
for _affinity, _overlap, _distance, tid, det_index in candidates:
|
|
if tid in used_tracks or det_index in used_detections:
|
|
continue
|
|
track = self._tracks.get(tid)
|
|
if track is None:
|
|
continue
|
|
det = detections[det_index]
|
|
old_box = list(track.box)
|
|
dt = max(0.05, timestamp - track.last_update)
|
|
new_box = [float(x) for x in det["box"]]
|
|
instant_velocity = [(new_box[i] - old_box[i]) / dt for i in range(4)]
|
|
track.velocity = [track.velocity[i] * 0.5 + instant_velocity[i] * 0.5 for i in range(4)]
|
|
track.box = new_box
|
|
track.score = float(det.get("score", 0) or 0)
|
|
track.missed = 0
|
|
track.hits += 1
|
|
track.observed = True
|
|
track.last_seen = timestamp
|
|
track.last_update = timestamp
|
|
track.algorithm_id = det.get("algorithm_id")
|
|
track.algorithm_name = det.get("algorithm_name", "")
|
|
if track.score >= 0.6 or track.hits >= 2:
|
|
track.confirmed = True
|
|
used_tracks.add(tid)
|
|
used_detections.add(det_index)
|
|
|
|
new_tracks = []
|
|
for det_index, det in enumerate(detections or []):
|
|
if det_index in used_detections or len(det.get("box") or []) != 4:
|
|
continue
|
|
tid = self._next_id
|
|
self._next_id += 1
|
|
self._tracks[tid] = Track(tid, det, frame_index, timestamp)
|
|
new_tracks.append(tid)
|
|
used_tracks.add(tid)
|
|
|
|
ended = []
|
|
for tid, track in list(self._tracks.items()):
|
|
if not track.observed:
|
|
track.missed += 1
|
|
if not track.confirmed:
|
|
track.hits = 0 # 中等置信度必须连续两次命中
|
|
if timestamp - track.last_seen > max_age:
|
|
ended.append(tid)
|
|
del self._tracks[tid]
|
|
|
|
active = [t.as_dict() for t in self._tracks.values() if t.observed]
|
|
return active, ended, new_tracks, frame_index
|
|
|
|
def display_tracks(self, timestamp=None):
|
|
now = float(timestamp if timestamp is not None else time.time())
|
|
result = []
|
|
for track in self._tracks.values():
|
|
item = track.as_dict()
|
|
item["age_sec"] = max(0.0, now - track.last_seen)
|
|
item["opacity"] = 1.0 if track.observed else max(0.15, 1.0 - item["age_sec"] / self.retention_sec)
|
|
result.append(item)
|
|
return result
|
|
|
|
def all_active(self):
|
|
return list(self._tracks.values())
|
|
|
|
def reset(self):
|
|
self._tracks.clear()
|
|
self._next_id = 1
|
|
self._last_update_ts = None
|
|
self._recent_period = 1.0
|