Create tracking.py

tracking.py

@@ -0,0 +1,197 @@
+import numpy as np
+from typing import List, Optional, Tuple
+from scipy.optimize import linear_sum_assignment
+from collections import deque
+import uuid
+
+class Track:
+    """Simple track for a single dog"""
+    
+    def __init__(self, detection: Detection, track_id: Optional[int] = None):
+        """Initialize track from first detection"""
+        self.track_id = track_id if track_id else self._generate_id()
+        self.bbox = detection.bbox
+        self.detections = [detection]
+        self.confidence = detection.confidence
+        
+        # Track state
+        self.age = 1
+        self.time_since_update = 0
+        self.state = 'tentative'  # tentative -> confirmed -> deleted
+        self.hits = 1
+        
+        # Store center points for trajectory
+        cx = (self.bbox[0] + self.bbox[2]) / 2
+        cy = (self.bbox[1] + self.bbox[3]) / 2
+        self.trajectory = deque(maxlen=30)
+        self.trajectory.append((cx, cy))
+        
+    def _generate_id(self) -> int:
+        """Generate unique track ID"""
+        return int(uuid.uuid4().int % 100000)
+    
+    def predict(self):
+        """Simple prediction - just use last position"""
+        self.age += 1
+        self.time_since_update += 1
+        
+    def update(self, detection: Detection):
+        """Update track with new detection"""
+        self.bbox = detection.bbox
+        self.detections.append(detection)
+        self.confidence = detection.confidence
+        
+        self.hits += 1
+        self.time_since_update = 0
+        
+        # Update trajectory
+        cx = (self.bbox[0] + self.bbox[2]) / 2
+        cy = (self.bbox[1] + self.bbox[3]) / 2
+        self.trajectory.append((cx, cy))
+        
+        # Confirm track after 3 hits
+        if self.state == 'tentative' and self.hits >= 3:
+            self.state = 'confirmed'
+            
+        # Keep only recent detections to save memory
+        if len(self.detections) > 10:
+            self.detections = self.detections[-10:]
+    
+    def mark_missed(self):
+        """Mark track as missed in current frame"""
+        if self.state == 'confirmed' and self.time_since_update > 15:
+            self.state = 'deleted'
+            
+class SimpleTracker:
+    """
+    Simplified ByteTrack - IoU-based tracking
+    Robust and proven approach without complexity
+    """
+    
+    def __init__(self, 
+                 match_threshold: float = 0.5,
+                 track_buffer: int = 30):
+        """
+        Initialize tracker
+        
+        Args:
+            match_threshold: IoU threshold for matching (0.5 works well)
+            track_buffer: Frames to keep lost tracks
+        """
+        self.match_threshold = match_threshold
+        self.track_buffer = track_buffer
+        
+        self.tracks: List[Track] = []
+        self.track_id_count = 1
+        
+    def update(self, detections: List[Detection]) -> List[Track]:
+        """
+        Update tracks with new detections
+        
+        Args:
+            detections: List of detections from current frame
+            
+        Returns:
+            List of active tracks
+        """
+        # Predict existing tracks
+        for track in self.tracks:
+            track.predict()
+        
+        # Get active tracks
+        active_tracks = [t for t in self.tracks if t.state != 'deleted']
+        
+        if len(detections) > 0 and len(active_tracks) > 0:
+            # Calculate IoU matrix
+            iou_matrix = self._calculate_iou_matrix(active_tracks, detections)
+            
+            # Hungarian matching
+            matched, unmatched_tracks, unmatched_dets = self._associate(
+                iou_matrix, self.match_threshold
+            )
+            
+            # Update matched tracks
+            for t_idx, d_idx in matched:
+                active_tracks[t_idx].update(detections[d_idx])
+            
+            # Mark unmatched tracks as missed
+            for t_idx in unmatched_tracks:
+                active_tracks[t_idx].mark_missed()
+            
+            # Create new tracks for unmatched detections
+            for d_idx in unmatched_dets:
+                new_track = Track(detections[d_idx], self.track_id_count)
+                self.track_id_count += 1
+                self.tracks.append(new_track)
+                
+        elif len(detections) > 0:
+            # No existing tracks - create new ones
+            for detection in detections:
+                new_track = Track(detection, self.track_id_count)
+                self.track_id_count += 1
+                self.tracks.append(new_track)
+        else:
+            # No detections - mark all as missed
+            for track in active_tracks:
+                track.mark_missed()
+        
+        # Remove deleted tracks
+        self.tracks = [t for t in self.tracks if t.state != 'deleted']
+        
+        # Return confirmed tracks
+        return [t for t in self.tracks if t.state == 'confirmed']
+    
+    def _calculate_iou_matrix(self, tracks: List[Track], 
+                             detections: List[Detection]) -> np.ndarray:
+        """Calculate IoU between all tracks and detections"""
+        matrix = np.zeros((len(tracks), len(detections)))
+        
+        for t_idx, track in enumerate(tracks):
+            for d_idx, detection in enumerate(detections):
+                matrix[t_idx, d_idx] = self._iou(track.bbox, detection.bbox)
+                
+        return matrix
+    
+    def _iou(self, bbox1: List[float], bbox2: List[float]) -> float:
+        """Calculate Intersection over Union"""
+        x1 = max(bbox1[0], bbox2[0])
+        y1 = max(bbox1[1], bbox2[1])
+        x2 = min(bbox1[2], bbox2[2])
+        y2 = min(bbox1[3], bbox2[3])
+        
+        if x2 < x1 or y2 < y1:
+            return 0.0
+        
+        intersection = (x2 - x1) * (y2 - y1)
+        area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
+        area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
+        union = area1 + area2 - intersection
+        
+        return intersection / union if union > 0 else 0.0
+    
+    def _associate(self, iou_matrix: np.ndarray, 
+                  threshold: float) -> Tuple[List, List, List]:
+        """Hungarian algorithm for optimal assignment"""
+        matched_indices = []
+        
+        if iou_matrix.max() >= threshold:
+            # Convert to cost matrix
+            cost_matrix = 1 - iou_matrix
+            row_ind, col_ind = linear_sum_assignment(cost_matrix)
+            
+            for r, c in zip(row_ind, col_ind):
+                if iou_matrix[r, c] >= threshold:
+                    matched_indices.append([r, c])
+        
+        unmatched_tracks = []
+        unmatched_detections = []
+        
+        for t in range(iou_matrix.shape[0]):
+            if t not in [m[0] for m in matched_indices]:
+                unmatched_tracks.append(t)
+                
+        for d in range(iou_matrix.shape[1]):
+            if d not in [m[1] for m in matched_indices]:
+                unmatched_detections.append(d)
+                
+        return matched_indices, unmatched_tracks, unmatched_detections