Spaces:
Running on Zero
Running on Zero
| """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() | |