Spaces:
Sleeping
Sleeping
| """Multi-object tracking module with BoT-SORT and ByteTrack support.""" | |
| import numpy as np | |
| import supervision as sv | |
| from config.settings import TrackerConfig | |
| class Tracker: | |
| """Wraps BoT-SORT (via boxmot) or ByteTrack (via supervision) for | |
| multi-object tracking with persistent ID assignment. | |
| BoT-SORT provides: | |
| - ReID appearance embeddings for identity preservation | |
| - Camera motion compensation (GMC) for handling pans/zooms | |
| - Two-stage association (high + low confidence detections) | |
| ByteTrack provides: | |
| - Fast IoU-only association (no ReID) | |
| - Good baseline for comparison | |
| """ | |
| def __init__(self, config: TrackerConfig): | |
| self.config = config | |
| self.tracker_type = config.tracker_type | |
| if config.tracker_type == "botsort": | |
| try: | |
| self._init_botsort(config) | |
| except ImportError: | |
| print("Warning: boxmot not installed, falling back to bytetrack") | |
| self.tracker_type = "bytetrack" | |
| self._init_bytetrack(config) | |
| elif config.tracker_type == "bytetrack": | |
| self._init_bytetrack(config) | |
| else: | |
| raise ValueError(f"Unknown tracker type: {config.tracker_type}") | |
| def _init_botsort(self, config: TrackerConfig): | |
| """Initialize BoT-SORT tracker via boxmot.""" | |
| from boxmot import BotSort | |
| self.tracker = BotSort( | |
| reid_weights=config.reid_model, | |
| device="cpu", # ReID on CPU for MPS compatibility | |
| half=False, | |
| track_high_thresh=config.track_high_thresh, | |
| track_low_thresh=config.track_low_thresh, | |
| new_track_thresh=config.new_track_thresh, | |
| track_buffer=config.track_buffer, | |
| match_thresh=config.match_thresh, | |
| proximity_thresh=config.proximity_thresh, | |
| appearance_thresh=config.appearance_thresh, | |
| cmc_method=config.cmc_method, | |
| frame_rate=config.frame_rate, | |
| ) | |
| def _init_bytetrack(self, config: TrackerConfig): | |
| """Initialize ByteTrack tracker via supervision.""" | |
| self.tracker = sv.ByteTrack( | |
| track_activation_threshold=config.track_high_thresh, | |
| lost_track_buffer=config.track_buffer, | |
| minimum_matching_threshold=config.match_thresh, | |
| frame_rate=config.frame_rate, | |
| ) | |
| def update(self, detections: sv.Detections, frame: np.ndarray) -> sv.Detections: | |
| """Update tracker with new detections and return tracked objects. | |
| Args: | |
| detections: Detection results from the detector. | |
| frame: Current video frame (needed for BoT-SORT ReID extraction). | |
| Returns: | |
| sv.Detections with tracker_id field populated. | |
| """ | |
| if self.tracker_type == "botsort": | |
| return self._update_botsort(detections, frame) | |
| else: | |
| return self._update_bytetrack(detections) | |
| def _update_botsort(self, detections: sv.Detections, frame: np.ndarray) -> sv.Detections: | |
| """Update BoT-SORT tracker. | |
| boxmot expects: numpy array of shape (N, 6) = [x1, y1, x2, y2, conf, cls] | |
| boxmot returns: numpy array of shape (M, 8) = [x1, y1, x2, y2, id, conf, cls, idx] | |
| """ | |
| if len(detections) == 0: | |
| # Still update tracker to age existing tracks | |
| empty = np.empty((0, 6)) | |
| self.tracker.update(empty, frame) | |
| return sv.Detections.empty() | |
| # Convert sv.Detections -> boxmot input format | |
| dets = np.column_stack([ | |
| detections.xyxy, | |
| detections.confidence, | |
| detections.class_id, | |
| ]) | |
| # Run tracker update | |
| tracked = self.tracker.update(dets, frame) | |
| if len(tracked) == 0: | |
| return sv.Detections.empty() | |
| # Convert boxmot output -> sv.Detections | |
| return sv.Detections( | |
| xyxy=tracked[:, 0:4].astype(np.float32), | |
| confidence=tracked[:, 5].astype(np.float32), | |
| class_id=tracked[:, 6].astype(int), | |
| tracker_id=tracked[:, 4].astype(int), | |
| ) | |
| def _update_bytetrack(self, detections: sv.Detections) -> sv.Detections: | |
| """Update ByteTrack tracker.""" | |
| if len(detections) == 0: | |
| return sv.Detections.empty() | |
| return self.tracker.update_with_detections(detections) | |
| def reset(self): | |
| """Reset tracker state (for new video or comparison runs).""" | |
| config = self.config | |
| if self.tracker_type == "botsort": | |
| self._init_botsort(config) | |
| else: | |
| self._init_bytetrack(config) | |