product-recognition-demo/recognization/code/tracking.py

"""
Session-scoped multi-object tracking for live recognition.

Wraps ultralytics' BYTETracker (Kalman prediction + two-stage association) with
one critical fix: ultralytics mints track IDs from a PROCESS-GLOBAL counter
(BaseTrack._count), and BYTETracker.__init__ resets that counter. With several
live sessions that means one client's tracker creation re-issues IDs already in
use by another client, and a cached label can land on the wrong product.

SessionTracker gives every tracker instance its own counter and never touches
the global one. The rest of the algorithm is used unmodified.

Pinned to ultralytics 8.4.153 (see requirements.txt): this module touches
tracker internals.
"""

from __future__ import annotations

import itertools
import os
import sys
import threading
from typing import Any, List, Optional, Sequence

import numpy as np

current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

from ultralytics.trackers.basetrack import BaseTrack, STrack
from ultralytics.trackers.byte_tracker import BYTETracker
from ultralytics.utils import IterableSimpleNamespace, yaml_load

_BYTETRACK_YAML = os.path.join(
    os.path.dirname(os.path.abspath(__import__("ultralytics").__file__)),
    "cfg", "trackers", "bytetrack.yaml",
)

DEFAULT_TRACKER_CONFIG = {
    "tracker_type": "bytetrack",
    "track_high_thresh": 0.25,
    "track_low_thresh": 0.1,
    "new_track_thresh": 0.25,
    "track_buffer": 30,
    "match_thresh": 0.8,
    "fuse_score": True,
}


def _default_args() -> IterableSimpleNamespace:
    try:
        return IterableSimpleNamespace(**yaml_load(_BYTETRACK_YAML))
    except Exception:
        return IterableSimpleNamespace(**DEFAULT_TRACKER_CONFIG)


class _Detections:
    """Duck-typed stand-in for ultralytics Results: only what BYTETracker reads.

    Must NOT expose `xywhr`; the tracker uses its presence to pick rotated-box
    handling.
    """

    __slots__ = ("xywh", "conf", "cls")

    def __init__(self, boxes_xyxy: np.ndarray, scores: np.ndarray):
        boxes = np.asarray(boxes_xyxy, dtype=np.float32).reshape(-1, 4)
        scores = np.asarray(scores, dtype=np.float32).reshape(-1)
        widths = boxes[:, 2] - boxes[:, 0]
        heights = boxes[:, 3] - boxes[:, 1]
        self.xywh = np.stack(
            [boxes[:, 0] + widths / 2.0, boxes[:, 1] + heights / 2.0, widths, heights], axis=1
        ).astype(np.float32)
        self.conf = scores
        self.cls = np.zeros(len(scores), dtype=np.float32)

    def __len__(self) -> int:
        return len(self.conf)

    def __getitem__(self, mask: Any) -> "_Detections":
        mask = np.asarray(mask)
        subset = _Detections.__new__(_Detections)
        subset.xywh = self.xywh[mask]
        subset.conf = self.conf[mask]
        subset.cls = self.cls[mask]
        return subset


class _SessionByteTrack(BYTETracker):
    """BYTETracker with a per-instance ID counter."""

    def __init__(self, args: IterableSimpleNamespace):
        self._id_counter = itertools.count(1)
        self._id_lock = threading.Lock()
        # super().__init__ calls self.reset_id(); the override below only touches
        # this instance's counter, so the global one is never disturbed.
        super().__init__(args)
        self.track_class = self._make_track_class()

    def _next_id(self) -> int:
        with self._id_lock:
            return next(self._id_counter)

    def _make_track_class(self):
        tracker = self

        class SessionSTrack(STrack):
            @staticmethod
            def next_id() -> int:
                return tracker._next_id()

        SessionSTrack.__name__ = "SessionSTrack"
        return SessionSTrack

    def reset_id(self) -> None:  # type: ignore[override]
        """Reset only this tracker's counter; never BaseTrack._count."""
        with self._id_lock:
            self._id_counter = itertools.count(1)
        self.track_class = self._make_track_class()


class TrackOut:
    """One tracked object for a frame."""

    __slots__ = ("track_id", "box", "score", "det_index")

    def __init__(self, track_id: int, box: List[int], score: float, det_index: Optional[int]):
        self.track_id = int(track_id)
        self.box = [int(v) for v in box]
        self.score = float(score)
        self.det_index = None if det_index is None or det_index < 0 else int(det_index)

    def __repr__(self) -> str:  # pragma: no cover - debug helper
        return f"TrackOut(id={self.track_id}, box={self.box}, score={self.score:.3f}, det={self.det_index})"


class SessionTracker:
    """Per-session tracker. Used by one session at a time, under its lock."""

    def __init__(
        self,
        *,
        high: float = 0.25,
        low: float = 0.1,
        new: float = 0.25,
        buffer: int = 30,
        match: float = 0.8,
        fuse: bool = True,
    ):
        args = _default_args()
        args.track_high_thresh = float(high)
        args.track_low_thresh = float(low)
        args.new_track_thresh = float(new)
        args.track_buffer = int(buffer)
        args.match_thresh = float(match)
        args.fuse_score = bool(fuse)
        self.args = args
        self._tracker = _SessionByteTrack(args)

    def update(
        self,
        boxes_xyxy: Sequence[Sequence[int]],
        scores: Sequence[float],
    ) -> List[TrackOut]:
        boxes = np.asarray(boxes_xyxy, dtype=np.float32).reshape(-1, 4)
        scores_array = np.asarray(scores, dtype=np.float32).reshape(-1)
        if len(boxes) != len(scores_array):
            raise ValueError("boxes and scores must have the same length")
        detections = _Detections(boxes, scores_array)
        rows = self._tracker.update(detections)
        outputs: List[TrackOut] = []
        if rows is None or len(rows) == 0:
            return outputs
        for row in np.asarray(rows):
            # [x1, y1, x2, y2, track_id, score, cls, idx]
            outputs.append(
                TrackOut(
                    track_id=int(row[4]),
                    box=[int(round(row[0])), int(round(row[1])), int(round(row[2])), int(round(row[3]))],
                    score=float(row[5]),
                    det_index=int(row[7]) if len(row) > 7 else None,
                )
            )
        return outputs

    def reset(self) -> None:
        self._tracker.reset()