"""Person detection and tracking (YOLO11 + ByteTrack); keeps the most prominent person.""" import queue import threading from collections import defaultdict import numpy as np import torch from tqdm import tqdm from ultralytics import YOLO from exmr.geometry.smoothing import moving_average_smooth from exmr.preprocessing.video import video_lwh class Tracker: """Tracks all people in a video and returns the boxes of the most prominent one. Args: checkpoint: YOLO11 detection weights (yolo11m.pt is fast, yolo11x.pt more accurate) """ def __init__(self, checkpoint): self.yolo = YOLO(checkpoint) # decode the video in a background thread while the detector runs self.yolo.add_callback("on_predict_start", _start_video_prefetch) self.yolo.add_callback("on_predict_end", _close_video_prefetch) def track(self, video_path): """Per frame, a list of {"id": track id, "bbx_xyxy": (4,) array} detections.""" results = self.yolo.track( video_path, device="cuda", conf=0.5, classes=0, # person verbose=False, stream=True, batch=8, half=True, persist=False, # ByteTrack instead of the default BoT-SORT, whose global motion compensation runs # on the CPU and dominates this stage, while we only need the main person's boxes tracker="bytetrack.yaml", ) # skip Ultralytics' warm-up: its 640x640 dummy image has a shape that the video's frames # never have, and cuDNN 9 spends ~1 s setting up each new input shape self.yolo.predictor.done_warmup = True track_history = [] try: for result in tqdm(results, total=video_lwh(video_path)[0], desc="Tracking"): frame = [] if result.boxes.id is not None: track_ids = result.boxes.id.int().cpu().tolist() bbx_xyxy = result.boxes.xyxy.cpu().numpy() # (N, 4) frame = [{"id": i, "bbx_xyxy": box} for i, box in zip(track_ids, bbx_xyxy)] track_history.append(frame) finally: try: results.close() finally: _close_video_prefetch(self.yolo.predictor) return track_history def get_main_track(self, video_path): """(L, 4) bounding boxes [x1, y1, x2, y2] of the most prominent person, i.e. the track with the largest summed box area; frames where it is missing are interpolated.""" L, w, h = video_lwh(video_path) id_to_frame_ids = defaultdict(list) id_to_bbx_xyxys = defaultdict(list) for frame_id, frame in enumerate(self.track(video_path)): for det in frame: id_to_frame_ids[det["id"]].append(frame_id) id_to_bbx_xyxys[det["id"]].append(det["bbx_xyxy"]) if not id_to_frame_ids: raise RuntimeError(f"No person detected in {video_path}") def area_sum(track_id): boxes = np.array(id_to_bbx_xyxys[track_id]) bbx_wh = boxes[:, 2:] - boxes[:, :2] return (bbx_wh[:, 0] * bbx_wh[:, 1] / w / h).sum() track_id = max(id_to_frame_ids, key=area_sum) # first one wins ties, like a stable sort frame_ids = torch.tensor(id_to_frame_ids[track_id]) bbx_xyxy = torch.tensor(np.array(id_to_bbx_xyxys[track_id])) # (N, 4) # fill missing frames by linear interpolation, then smooth present = torch.zeros(L, dtype=torch.bool) present[frame_ids] = True if len(frame_ids) < L: full = torch.zeros((L, 4), dtype=bbx_xyxy.dtype) full[present] = bbx_xyxy bbx_xyxy = _interpolate_missing(full, present) bbx_xyxy = moving_average_smooth(bbx_xyxy, window_size=5, dim=0) return moving_average_smooth(bbx_xyxy, window_size=5, dim=0) def _interpolate_missing(data, present): """Linear interpolation over each run of missing frames; runs at the start/end copy the nearest present frame.""" data = data.clone() for missing in _runs_of_true(~present): if missing[0] - 1 < 0: data[missing] = data[missing[-1] + 1].clone() elif missing[-1] + 1 >= len(data): data[missing] = data[missing[0] - 1].clone() else: prev = data[missing[0] - 1] next = data[missing[-1] + 1] t = torch.linspace(0, 1, len(missing) + 2)[1:-1][:, None] data[missing] = t * (next - prev)[None] + prev[None] return data def _runs_of_true(mask): """(L,) bool -> list of index tensors, one per run of consecutive True values.""" padded = torch.cat([torch.tensor([False]), mask, torch.tensor([False])]) diffs = torch.diff(padded.int()) starts = (diffs == 1).nonzero().flatten() ends = (diffs == -1).nonzero().flatten() return [torch.arange(s, e) for s, e in zip(starts, ends)] class _PrefetchVideo: """Wraps ultralytics' video loader to decode in a background thread (at most two batches ahead). Only the worker thread touches the OpenCV capture; batches are consumed in order.""" def __init__(self, loader): self.loader = loader self._pending = queue.Queue(maxsize=2) self._stop = threading.Event() self._thread = None def __getattr__(self, name): return getattr(self.loader, name) def __len__(self): return len(self.loader) def _put(self, item): while not self._stop.is_set(): try: self._pending.put(item, timeout=0.1) return except queue.Full: pass def _decode(self): try: for batch in self.loader: if self._stop.is_set(): break self._put((True, batch)) except Exception as error: self._put((False, error)) finally: try: if self.loader.cap is not None: self.loader.cap.release() finally: self._put((True, None)) def __iter__(self): if self._thread is None and not self._stop.is_set(): self._thread = threading.Thread(target=self._decode, daemon=True) self._thread.start() return self def __next__(self): if self._stop.is_set(): raise StopIteration if self._thread is None: iter(self) ok, item = self._pending.get() if not ok: self.close() raise item if item is None: self.close() raise StopIteration return item def close(self): self._stop.set() if self._thread is not None: self._thread.join(timeout=5) if self._thread.is_alive(): raise RuntimeError("Video decoder did not stop within five seconds") elif self.loader.cap is not None: self.loader.cap.release() while True: # release queued frames even if inference exited early try: self._pending.get_nowait() except queue.Empty: break def _start_video_prefetch(predictor): loader = predictor.dataset if loader.nf != 1 or loader.mode != "video": raise ValueError("Tracker decode prefetch expects exactly one video file") predictor.dataset = _PrefetchVideo(loader) def _close_video_prefetch(predictor): if isinstance(getattr(predictor, "dataset", None), _PrefetchVideo): predictor.dataset.close()