"""Isolated adapter for Ultralytics' ByteTrack implementation.""" from __future__ import annotations import logging import inspect from argparse import Namespace from typing import Any import numpy as np from .config_loader import TrackerSettings from .exceptions import TrackerError from .models import BoundingBox, Detection, FrameInfo, TrackedPerson LOGGER = logging.getLogger(__name__) class _TrackerInput: """Minimal Results-like input expected by current Ultralytics ByteTrack.""" def __init__(self, detections: list[Detection]) -> None: self.xyxy = np.asarray([[d.bbox.x1, d.bbox.y1, d.bbox.x2, d.bbox.y2] for d in detections], dtype=np.float32).reshape(-1, 4) self.conf = np.asarray([d.confidence for d in detections], dtype=np.float32) self.cls = np.asarray([d.class_id for d in detections], dtype=np.float32) @property def xywh(self) -> np.ndarray: values = self.xyxy.copy() values[:, 0] = (self.xyxy[:, 0] + self.xyxy[:, 2]) / 2 values[:, 1] = (self.xyxy[:, 1] + self.xyxy[:, 3]) / 2 values[:, 2] = self.xyxy[:, 2] - self.xyxy[:, 0] values[:, 3] = self.xyxy[:, 3] - self.xyxy[:, 1] return values def __len__(self) -> int: return len(self.conf) def __getitem__(self, index: Any) -> "_TrackerInput": """Support boolean/slice indexing used by newer Ultralytics releases.""" selected = object.__new__(_TrackerInput) selected.xyxy = self.xyxy[index] selected.conf = self.conf[index] selected.cls = self.cls[index] if selected.xyxy.ndim == 1: selected.xyxy = selected.xyxy.reshape(1, 4) selected.conf = selected.conf.reshape(1) selected.cls = selected.cls.reshape(1) return selected class PersonTracker: """Convert domain detections to ByteTrack and normalize its output.""" def __init__(self, config: TrackerSettings, frame_rate: float = 30.0) -> None: self.config = config self._frame_rate = frame_rate if frame_rate > 0 else 30.0 self._tracker: Any = None self._raw_to_public: dict[int, int] = {} self._next_public_id = 1 self._ages: dict[int, int] = {} self._first_frames: dict[int, int] = {} if config.enabled: self._initialize() def _initialize(self) -> None: try: from ultralytics.trackers.byte_tracker import BYTETracker args = Namespace( tracker_type="bytetrack", track_high_thresh=self.config.track_high_thresh, track_low_thresh=self.config.track_low_thresh, new_track_thresh=self.config.new_track_thresh, match_thresh=self.config.match_thresh, track_buffer=self.config.track_buffer_frames, fuse_score=self.config.fuse_score, ) self._tracker = self._create_backend(BYTETracker, args, round(self._frame_rate)) except Exception as exc: raise TrackerError(f"Failed to initialize Ultralytics ByteTrack: {exc}") from exc @staticmethod def _create_backend(backend_class: type[Any], args: Namespace, frame_rate: int) -> Any: """Initialize ByteTrack across old and new Ultralytics constructor APIs.""" try: parameters = inspect.signature(backend_class.__init__).parameters except (TypeError, ValueError): parameters = {} if "frame_rate" in parameters: return backend_class(args=args, frame_rate=frame_rate) return backend_class(args) def update(self, detections: list[Detection], frame: np.ndarray, frame_info: FrameInfo) -> list[TrackedPerson]: """Update ByteTrack; a one-frame failure is logged and returned as empty.""" if not self.config.enabled: return [] if self._tracker is None: raise TrackerError("Tracker is not initialized") try: raw_tracks = self._tracker.update(_TrackerInput([d for d in detections if d.class_id == 0]), frame) return self._normalize(raw_tracks, frame.shape[1], frame.shape[0], detections, frame_info.frame_number) except Exception: LOGGER.exception("ByteTrack update failed for one frame") return [] def _normalize( self, raw_tracks: Any, width: int, height: int, detections: list[Detection] | None = None, frame_number: int | None = None, ) -> list[TrackedPerson]: persons: list[TrackedPerson] = [] for raw in raw_tracks if raw_tracks is not None else []: row = raw.tolist() if hasattr(raw, "tolist") else list(raw) if len(row) < 7: continue x1, y1, x2, y2 = (round(float(value)) for value in row[:4]) x1, y1 = max(0, min(width - 1, x1)), max(0, min(height - 1, y1)) x2, y2 = max(0, min(width - 1, x2)), max(0, min(height - 1, y2)) if x1 >= x2 or y1 >= y2: continue raw_id, confidence, class_id = int(row[4]), float(row[5]), int(row[6]) if class_id != 0: continue if raw_id not in self._raw_to_public: self._raw_to_public[raw_id] = self._allocate_id() public_id = self._raw_to_public[raw_id] if frame_number is not None: self._first_frames.setdefault(public_id, frame_number) self._ages[public_id] = frame_number - self._first_frames[public_id] + 1 else: self._ages[public_id] = self._ages.get(public_id, 0) + 1 bbox = BoundingBox(x1, y1, x2, y2) detection_index = int(row[7]) if len(row) > 7 else -1 source_detection = detections[detection_index] if detections and 0 <= detection_index < len(detections) else None persons.append(TrackedPerson( track_id=public_id, bbox=bbox, confidence=max(0.0, min(1.0, confidence)), class_id=class_id, class_name="person", center_x=float(bbox.center_x), center_y=float(bbox.center_y), foot_x=float(bbox.foot_x), foot_y=float(bbox.foot_y), inside_evaluation_region=source_detection.inside_evaluation_region if source_detection else False, is_confirmed=self._ages[public_id] >= self.config.min_confirmed_frames, track_age_frames=self._ages[public_id], missed_frames=0, )) return persons def _allocate_id(self) -> int: value = self._next_public_id self._next_public_id += 1 return value def reset(self) -> None: """Reset tracker state while keeping public IDs monotonic for this run.""" self._raw_to_public.clear() self._ages.clear() self._first_frames.clear() if self.config.enabled: self._initialize() def close(self) -> None: self._tracker = None self._raw_to_public.clear() self._ages.clear() self._first_frames.clear() @property def tracker_name(self) -> str: return "ByteTrack (Ultralytics)" if self.config.enabled else "disabled"