video_monitor/app/analysis/tracker.py

"""低帧率视频目标跟踪。

采用同类别全局一对一匹配，并用时间而不是丢帧数保留轨迹。这样在 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