tracker_module.py

import time
from dataclasses import dataclass


def iou_xyxy(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.0, ix2 - ix1)
    ih = max(0.0, iy2 - iy1)
    inter = iw * ih
    if inter <= 0:
        return 0.0
    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
    if union <= 0:
        return 0.0
    return inter / union


@dataclass
class Track:
    track_id: int
    bbox: tuple
    cls_name: str
    conf: float
    first_seen: float
    last_seen: float
    hits: int = 1
    miss: int = 0


class SORTTracker:
    """A lightweight SORT-style IoU tracker for CPU environments."""

    def __init__(self, iou_threshold=0.3, max_miss=15):
        self.iou_threshold = float(iou_threshold)
        self.max_miss = int(max_miss)
        self._tracks = {}
        self._next_id = 1

    def update(self, detections):
        """
        detections: [{"bbox": (x1,y1,x2,y2), "class": str, "confidence": float}]
        returns list of active tracks with stable IDs.
        """
        now = time.time()
        dets = list(detections or [])
        matched_tracks = set()
        matched_dets = set()

        track_ids = list(self._tracks.keys())
        scored_pairs = []
        for tid in track_ids:
            tr = self._tracks[tid]
            for di, det in enumerate(dets):
                score = iou_xyxy(tr.bbox, det["bbox"])
                if score >= self.iou_threshold:
                    scored_pairs.append((score, tid, di))

        scored_pairs.sort(reverse=True, key=lambda x: x[0])
        for score, tid, di in scored_pairs:
            if tid in matched_tracks or di in matched_dets:
                continue
            det = dets[di]
            tr = self._tracks[tid]
            tr.bbox = tuple(det["bbox"])
            tr.cls_name = det["class"]
            tr.conf = float(det["confidence"])
            tr.last_seen = now
            tr.hits += 1
            tr.miss = 0
            matched_tracks.add(tid)
            matched_dets.add(di)

        for di, det in enumerate(dets):
            if di in matched_dets:
                continue
            tid = self._next_id
            self._next_id += 1
            self._tracks[tid] = Track(
                track_id=tid,
                bbox=tuple(det["bbox"]),
                cls_name=det["class"],
                conf=float(det["confidence"]),
                first_seen=now,
                last_seen=now,
            )

        to_delete = []
        for tid, tr in self._tracks.items():
            if tid not in matched_tracks:
                tr.miss += 1
            if tr.miss > self.max_miss:
                to_delete.append(tid)
        for tid in to_delete:
            del self._tracks[tid]

        output = []
        for tid, tr in self._tracks.items():
            output.append(
                {
                    "track_id": tid,
                    "bbox": tr.bbox,
                    "class": tr.cls_name,
                    "confidence": tr.conf,
                    "age_sec": max(0.0, now - tr.first_seen),
                    "hits": tr.hits,
                }
            )
        return output