167 lines
7.0 KiB
Python
167 lines
7.0 KiB
Python
"""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"
|