AI-Talent-Force/dev_caio

1"""2ShortSmith v2 - Object Tracker Module3 4Multi-object tracking using ByteTrack for:5- Maintaining person identity across frames6- Handling occlusions and reappearances7- Tracking specific individuals through video8 9ByteTrack uses two-stage association for robust tracking.10"""11 12from pathlib import Path13from typing import List, Optional, Dict, Tuple, Union14from dataclasses import dataclass, field15import numpy as np16 17from utils.logger import get_logger, LogTimer18from utils.helpers import InferenceError19from config import get_config20 21logger = get_logger("models.tracker")22 23 24@dataclass25class TrackedObject:26    """Represents a tracked object across frames."""27    track_id: int                         # Unique track identifier28    bbox: Tuple[int, int, int, int]       # Current bounding box (x1, y1, x2, y2)29    confidence: float                      # Detection confidence30    class_id: int = 0                      # Object class (0 = person)31    frame_id: int = 0                      # Current frame number32 33    # Track history34    history: List[Tuple[int, int, int, int]] = field(default_factory=list)35    age: int = 0                          # Frames since first detection36    hits: int = 0                         # Number of detections37    time_since_update: int = 0            # Frames since last detection38 39    @property40    def center(self) -> Tuple[int, int]:41        x1, y1, x2, y2 = self.bbox42        return ((x1 + x2) // 2, (y1 + y2) // 2)43 44    @property45    def area(self) -> int:46        x1, y1, x2, y2 = self.bbox47        return (x2 - x1) * (y2 - y1)48 49    @property50    def is_confirmed(self) -> bool:51        """Track is confirmed after multiple detections."""52        return self.hits >= 353 54 55@dataclass56class TrackingResult:57    """Result of tracking for a single frame."""58    frame_id: int59    tracks: List[TrackedObject]60    lost_tracks: List[int]  # Track IDs lost this frame61    new_tracks: List[int]   # New track IDs this frame62 63 64class ObjectTracker:65    """66    Multi-object tracker using ByteTrack algorithm.67 68    ByteTrack features:69    - Two-stage association (high-confidence first, then low-confidence)70    - Handles occlusions by keeping lost tracks71    - Re-identifies objects after temporary disappearance72    """73 74    def __init__(75        self,76        track_thresh: float = 0.5,77        track_buffer: int = 30,78        match_thresh: float = 0.8,79    ):80        """81        Initialize tracker.82 83        Args:84            track_thresh: Detection confidence threshold for new tracks85            track_buffer: Frames to keep lost tracks86            match_thresh: IoU threshold for matching87        """88        self.track_thresh = track_thresh89        self.track_buffer = track_buffer90        self.match_thresh = match_thresh91 92        self._tracks: Dict[int, TrackedObject] = {}93        self._lost_tracks: Dict[int, TrackedObject] = {}94        self._next_id = 195        self._frame_id = 096 97        logger.info(98            f"ObjectTracker initialized (thresh={track_thresh}, "99            f"buffer={track_buffer}, match={match_thresh})"100        )101 102    def update(103        self,104        detections: List[Tuple[Tuple[int, int, int, int], float]],105    ) -> TrackingResult:106        """107        Update tracker with new detections.108 109        Args:110            detections: List of (bbox, confidence) tuples111 112        Returns:113            TrackingResult with current tracks114        """115        self._frame_id += 1116 117        if not detections:118            # No detections - age all tracks119            return self._handle_no_detections()120 121        # Separate high and low confidence detections122        high_conf = [(bbox, conf) for bbox, conf in detections if conf >= self.track_thresh]123        low_conf = [(bbox, conf) for bbox, conf in detections if conf < self.track_thresh]124 125        # First association: match high-confidence detections to active tracks126        matched, unmatched_tracks, unmatched_dets = self._associate(127            list(self._tracks.values()),128            high_conf,129            self.match_thresh,130        )131 132        # Update matched tracks133        for track_id, det_idx in matched:134            bbox, conf = high_conf[det_idx]135            self._update_track(track_id, bbox, conf)136 137        # Second association: match low-confidence to remaining tracks138        if low_conf and unmatched_tracks:139            remaining_tracks = [self._tracks[tid] for tid in unmatched_tracks]140            matched2, unmatched_tracks, _ = self._associate(141                remaining_tracks,142                low_conf,143                self.match_thresh * 0.9,  # Lower threshold144            )145 146            for track_id, det_idx in matched2:147                bbox, conf = low_conf[det_idx]148                self._update_track(track_id, bbox, conf)149 150        # Handle unmatched tracks151        lost_this_frame = []152        for track_id in unmatched_tracks:153            track = self._tracks[track_id]154            track.time_since_update += 1155 156            if track.time_since_update > self.track_buffer:157                # Remove track158                del self._tracks[track_id]159                lost_this_frame.append(track_id)160            else:161                # Move to lost tracks162                self._lost_tracks[track_id] = self._tracks.pop(track_id)163 164        # Try to recover lost tracks with unmatched detections165        recovered = self._recover_lost_tracks(166            [(high_conf[i] if i < len(high_conf) else low_conf[i - len(high_conf)])167             for i in unmatched_dets]168        )169 170        # Create new tracks for remaining detections171        new_tracks = []172        for i in unmatched_dets:173            if i not in recovered:174                det = high_conf[i] if i < len(high_conf) else low_conf[i - len(high_conf)]175                bbox, conf = det176                track_id = self._create_track(bbox, conf)177                new_tracks.append(track_id)178 179        return TrackingResult(180            frame_id=self._frame_id,181            tracks=list(self._tracks.values()),182            lost_tracks=lost_this_frame,183            new_tracks=new_tracks,184        )185 186    def _associate(187        self,188        tracks: List[TrackedObject],189        detections: List[Tuple[Tuple[int, int, int, int], float]],190        thresh: float,191    ) -> Tuple[List[Tuple[int, int]], List[int], List[int]]:192        """193        Associate detections to tracks using IoU.194 195        Returns:196            (matched_pairs, unmatched_track_ids, unmatched_detection_indices)197        """198        if not tracks or not detections:199            return [], [t.track_id for t in tracks], list(range(len(detections)))200 201        # Compute IoU matrix202        iou_matrix = np.zeros((len(tracks), len(detections)))203 204        for i, track in enumerate(tracks):205            for j, (det_bbox, _) in enumerate(detections):206                iou_matrix[i, j] = self._compute_iou(track.bbox, det_bbox)207 208        # Greedy matching209        matched = []210        unmatched_tracks = set(t.track_id for t in tracks)211        unmatched_dets = set(range(len(detections)))212 213        while True:214            # Find best match215            if iou_matrix.size == 0:216                break217 218            max_iou = np.max(iou_matrix)219            if max_iou < thresh:220                break221 222            max_idx = np.unravel_index(np.argmax(iou_matrix), iou_matrix.shape)223            track_idx, det_idx = max_idx224 225            track_id = tracks[track_idx].track_id226            matched.append((track_id, det_idx))227            unmatched_tracks.discard(track_id)228            unmatched_dets.discard(det_idx)229 230            # Remove matched row and column231            iou_matrix[track_idx, :] = -1232            iou_matrix[:, det_idx] = -1233 234        return matched, list(unmatched_tracks), list(unmatched_dets)235 236    def _compute_iou(237        self,238        bbox1: Tuple[int, int, int, int],239        bbox2: Tuple[int, int, int, int],240    ) -> float:241        """Compute IoU between two bounding boxes."""242        x1_1, y1_1, x2_1, y2_1 = bbox1243        x1_2, y1_2, x2_2, y2_2 = bbox2244 245        # Intersection246        xi1 = max(x1_1, x1_2)247        yi1 = max(y1_1, y1_2)248        xi2 = min(x2_1, x2_2)249        yi2 = min(y2_1, y2_2)250 251        if xi2 <= xi1 or yi2 <= yi1:252            return 0.0253 254        intersection = (xi2 - xi1) * (yi2 - yi1)255 256        # Union257        area1 = (x2_1 - x1_1) * (y2_1 - y1_1)258        area2 = (x2_2 - x1_2) * (y2_2 - y1_2)259        union = area1 + area2 - intersection260 261        return intersection / union if union > 0 else 0.0262 263    def _update_track(264        self,265        track_id: int,266        bbox: Tuple[int, int, int, int],267        confidence: float,268    ) -> None:269        """Update an existing track."""270        track = self._tracks.get(track_id) or self._lost_tracks.get(track_id)271 272        if track is None:273            return274 275        # Move from lost to active if needed276        if track_id in self._lost_tracks:277            self._tracks[track_id] = self._lost_tracks.pop(track_id)278 279        track = self._tracks[track_id]280        track.history.append(track.bbox)281        track.bbox = bbox282        track.confidence = confidence283        track.frame_id = self._frame_id284        track.hits += 1285        track.time_since_update = 0286 287    def _create_track(288        self,289        bbox: Tuple[int, int, int, int],290        confidence: float,291    ) -> int:292        """Create a new track."""293        track_id = self._next_id294        self._next_id += 1295 296        track = TrackedObject(297            track_id=track_id,298            bbox=bbox,299            confidence=confidence,300            frame_id=self._frame_id,301            age=1,302            hits=1,303        )304 305        self._tracks[track_id] = track306        logger.debug(f"Created new track {track_id}")307        return track_id308 309    def _recover_lost_tracks(310        self,311        detections: List[Tuple[Tuple[int, int, int, int], float]],312    ) -> set:313        """Try to recover lost tracks with unmatched detections."""314        recovered = set()315 316        if not self._lost_tracks or not detections:317            return recovered318 319        for det_idx, (bbox, conf) in enumerate(detections):320            best_iou = 0321            best_track_id = None322 323            for track_id, track in self._lost_tracks.items():324                iou = self._compute_iou(track.bbox, bbox)325                if iou > best_iou and iou > self.match_thresh * 0.7:326                    best_iou = iou327                    best_track_id = track_id328 329            if best_track_id is not None:330                self._update_track(best_track_id, bbox, conf)331                recovered.add(det_idx)332                logger.debug(f"Recovered track {best_track_id}")333 334        return recovered335 336    def _handle_no_detections(self) -> TrackingResult:337        """Handle frame with no detections."""338        lost_this_frame = []339 340        for track_id in list(self._tracks.keys()):341            track = self._tracks[track_id]342            track.time_since_update += 1343 344            if track.time_since_update > self.track_buffer:345                del self._tracks[track_id]346                lost_this_frame.append(track_id)347            else:348                self._lost_tracks[track_id] = self._tracks.pop(track_id)349 350        return TrackingResult(351            frame_id=self._frame_id,352            tracks=list(self._tracks.values()),353            lost_tracks=lost_this_frame,354            new_tracks=[],355        )356 357    def get_track(self, track_id: int) -> Optional[TrackedObject]:358        """Get a specific track by ID."""359        return self._tracks.get(track_id) or self._lost_tracks.get(track_id)360 361    def get_active_tracks(self) -> List[TrackedObject]:362        """Get all active tracks."""363        return list(self._tracks.values())364 365    def get_confirmed_tracks(self) -> List[TrackedObject]:366        """Get only confirmed tracks (multiple detections)."""367        return [t for t in self._tracks.values() if t.is_confirmed]368 369    def reset(self) -> None:370        """Reset tracker state."""371        self._tracks.clear()372        self._lost_tracks.clear()373        self._frame_id = 0374        logger.info("Tracker reset")375 376    def get_track_for_target(377        self,378        target_bbox: Tuple[int, int, int, int],379        threshold: float = 0.5,380    ) -> Optional[int]:381        """382        Find track that best matches a target bounding box.383 384        Args:385            target_bbox: Target bounding box to match386            threshold: Minimum IoU for match387 388        Returns:389            Track ID if found, None otherwise390        """391        best_iou = 0392        best_track = None393 394        for track in self._tracks.values():395            iou = self._compute_iou(track.bbox, target_bbox)396            if iou > best_iou and iou > threshold:397                best_iou = iou398                best_track = track.track_id399 400        return best_track401 402 403# Export public interface404__all__ = ["ObjectTracker", "TrackedObject", "TrackingResult"]405