Phase 11

src/tracker.py

@@ -0,0 +1,166 @@
+"""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"