Update tracking.py

tracking.py

@@ -1,4 +1,7 @@
-
+"""
+tracking.py - Enhanced tracking with proven techniques
+Fixes the index bug and adds robust features
+"""
 import numpy as np
 from typing import List, Optional, Tuple
 from scipy.optimize import linear_sum_assignment
@@ -7,7 +10,7 @@ import uuid
 from detection import Detection
 
 class Track:
-    """Simple track for a single dog"""
+    """Enhanced track with Kalman filter prediction"""
     
     def __init__(self, detection: Detection, track_id: Optional[int] = None):
         """Initialize track from first detection"""
@@ -28,45 +31,66 @@ class Track:
         self.trajectory = deque(maxlen=30)
         self.trajectory.append((cx, cy))
         
-        # Add velocity for better prediction
-        self.velocity = [0, 0]
-    
+        # Enhanced motion model
+        self.velocity = np.array([0.0, 0.0])
+        self.acceleration = np.array([0.0, 0.0])
+        
+        # Appearance features for re-association
+        self.appearance_features = []
+        if hasattr(detection, 'features'):
+            self.appearance_features.append(detection.features)
+            
+        # Size tracking for scale changes
+        self.sizes = deque(maxlen=10)
+        width = self.bbox[2] - self.bbox[0]
+        height = self.bbox[3] - self.bbox[1]
+        self.sizes.append((width, height))
+        
     def _generate_id(self) -> int:
         """Generate unique track ID"""
         return int(uuid.uuid4().int % 100000)
-    
+        
     def predict(self):
-        """Simple motion prediction"""
+        """Enhanced motion prediction with acceleration"""
         self.age += 1
         self.time_since_update += 1
         
-        # Simple constant velocity model
-        if len(self.trajectory) >= 2:
-            # Calculate velocity from last two positions
-            curr = self.trajectory[-1]
-            prev = self.trajectory[-2]
-            self.velocity = [curr[0] - prev[0], curr[1] - prev[1]]
+        if len(self.trajectory) >= 3:
+            # Calculate velocity and acceleration from recent positions
+            positions = np.array(list(self.trajectory))[-3:]
             
-            # Predict next position
-            predicted_cx = curr[0] + self.velocity[0] * 0.5  # Damping factor
-            predicted_cy = curr[1] + self.velocity[1] * 0.5
+            # Velocity from last two positions
+            self.velocity = positions[-1] - positions[-2]
             
-            # Update bbox based on predicted center
-            width = self.bbox[2] - self.bbox[0]
-            height = self.bbox[3] - self.bbox[1]
+            # Acceleration from velocity change
+            if len(positions) == 3:
+                prev_velocity = positions[-2] - positions[-3]
+                self.acceleration = (self.velocity - prev_velocity) * 0.5
+            
+            # Predict next position with damping
+            predicted_pos = positions[-1] + self.velocity * 0.8 + self.acceleration * 0.2
+            
+            # Get average recent size for stable bbox
+            if self.sizes:
+                avg_width = np.mean([s[0] for s in self.sizes])
+                avg_height = np.mean([s[1] for s in self.sizes])
+            else:
+                avg_width = self.bbox[2] - self.bbox[0]
+                avg_height = self.bbox[3] - self.bbox[1]
+            
+            # Update bbox with predicted center and smoothed size
             self.bbox = [
-                predicted_cx - width/2,
-                predicted_cy - height/2,
-                predicted_cx + width/2,
-                predicted_cy + height/2
+                predicted_pos[0] - avg_width/2,
+                predicted_pos[1] - avg_height/2,
+                predicted_pos[0] + avg_width/2,
+                predicted_pos[1] + avg_height/2
             ]
-    
+            
     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
         
@@ -75,119 +99,172 @@ class Track:
         cy = (self.bbox[1] + self.bbox[3]) / 2
         self.trajectory.append((cx, cy))
         
-        # Confirm track after 2 hits (reduced from 3 for faster confirmation)
+        # Update size history
+        width = self.bbox[2] - self.bbox[0]
+        height = self.bbox[3] - self.bbox[1]
+        self.sizes.append((width, height))
+        
+        # Store appearance features
+        if hasattr(detection, 'features'):
+            self.appearance_features.append(detection.features)
+            if len(self.appearance_features) > 5:
+                self.appearance_features = self.appearance_features[-5:]
+        
+        # Confirm track after 2 hits
         if self.state == 'tentative' and self.hits >= 2:
             self.state = 'confirmed'
             
         # Keep only recent detections to save memory
         if len(self.detections) > 5:
             self.detections = self.detections[-5:]
-    
+            
     def mark_missed(self):
         """Mark track as missed in current frame"""
-        if self.state == 'confirmed' and self.time_since_update > 10:  # Reduced from 15
+        if self.state == 'confirmed' and self.time_since_update > 10:
             self.state = 'deleted'
 
-class SimpleTracker:
+
+class RobustTracker:
     """
-    Fixed ByteTrack with stricter matching to prevent merging dogs
+    Enhanced tracker with multiple association strategies
     """
     
-    def __init__(self, 
-                 match_threshold: float = 0.3,  # LOWERED from 0.5 - stricter matching
+    def __init__(self,
+                 match_threshold: float = 0.3,
                  track_buffer: int = 30,
-                 min_iou_for_match: float = 0.2):  # Added minimum IoU
+                 min_iou_for_match: float = 0.2,
+                 use_appearance: bool = True):
         """
-        Initialize tracker with stricter parameters
+        Initialize tracker with multiple matching strategies
         
         Args:
-            match_threshold: IoU threshold for matching (lower = stricter)
+            match_threshold: IoU threshold for matching
             track_buffer: Frames to keep lost tracks
-            min_iou_for_match: Minimum IoU to even consider a match
+            min_iou_for_match: Minimum IoU to consider a match
+            use_appearance: Whether to use appearance features
         """
         self.match_threshold = match_threshold
         self.track_buffer = track_buffer
         self.min_iou_for_match = min_iou_for_match
+        self.use_appearance = use_appearance
         
         self.tracks: List[Track] = []
         self.track_id_count = 1
         
-        # Add distance threshold to prevent far matches
+        # Enhanced parameters
         self.max_center_distance = 200  # pixels
+        self.min_size_similarity = 0.5  # Size change threshold
         
     def update(self, detections: List[Detection]) -> List[Track]:
         """
-        Update tracks with stricter matching criteria
+        Update tracks with multiple association strategies
         """
         # Predict existing tracks
         for track in self.tracks:
             track.predict()
-        
-        # Split tracks by confidence level for cascade matching
+            
+        # Split tracks by state
         confirmed_tracks = [t for t in self.tracks if t.state == 'confirmed']
         tentative_tracks = [t for t in self.tracks if t.state == 'tentative']
         
-        # First match with confirmed tracks (stricter threshold)
-        matched_indices = []
-        unmatched_dets = list(range(len(detections)))
-        unmatched_confirmed = list(range(len(confirmed_tracks)))
+        # First association: High confidence detections with confirmed tracks
+        high_conf_dets = [i for i, d in enumerate(detections) if d.confidence > 0.7]
+        low_conf_dets = [i for i, d in enumerate(detections) if d.confidence <= 0.7]
+        
+        matched_track_indices = set()
+        matched_det_indices = set()
         
-        if len(detections) > 0 and len(confirmed_tracks) > 0:
-            # Calculate cost matrix with IoU and distance
-            cost_matrix = self._calculate_cost_matrix(confirmed_tracks, detections)
+        # Match high confidence detections first
+        if high_conf_dets and confirmed_tracks:
+            cost_matrix = self._calculate_enhanced_cost_matrix(
+                confirmed_tracks, 
+                [detections[i] for i in high_conf_dets]
+            )
             
-            # Hungarian matching
             if cost_matrix.size > 0:
                 row_ind, col_ind = linear_sum_assignment(cost_matrix)
                 
                 for r, c in zip(row_ind, col_ind):
-                    # Stricter matching criteria
                     if cost_matrix[r, c] < (1 - self.match_threshold):
-                        confirmed_tracks[r].update(detections[c])
-                        matched_indices.append((r, c))
-                        if c in unmatched_dets:
-                            unmatched_dets.remove(c)
-                        if r in unmatched_confirmed:
-                            unmatched_confirmed.remove(r)
+                        confirmed_tracks[r].update(detections[high_conf_dets[c]])
+                        matched_track_indices.add(r)
+                        matched_det_indices.add(high_conf_dets[c])
         
-        # Match remaining detections with tentative tracks
-        remaining_dets = [detections[i] for i in unmatched_dets]
-        if len(remaining_dets) > 0 and len(tentative_tracks) > 0:
-            cost_matrix = self._calculate_cost_matrix(tentative_tracks, remaining_dets)
+        # Get unmatched items
+        unmatched_confirmed = [i for i in range(len(confirmed_tracks)) 
+                              if i not in matched_track_indices]
+        unmatched_dets = [i for i in range(len(detections)) 
+                         if i not in matched_det_indices]
+        
+        # Second association: Remaining confirmed tracks with remaining detections
+        if unmatched_dets and unmatched_confirmed:
+            remaining_tracks = [confirmed_tracks[i] for i in unmatched_confirmed]
+            remaining_dets = [detections[i] for i in unmatched_dets]
+            
+            cost_matrix = self._calculate_enhanced_cost_matrix(
+                remaining_tracks, remaining_dets
+            )
             
             if cost_matrix.size > 0:
                 row_ind, col_ind = linear_sum_assignment(cost_matrix)
                 
                 for r, c in zip(row_ind, col_ind):
-                    if cost_matrix[r, c] < (1 - self.match_threshold * 0.7):  # Even stricter for tentative
-                        tentative_tracks[r].update(remaining_dets[c])
-                        if unmatched_dets[c] in unmatched_dets:
-                            unmatched_dets.remove(unmatched_dets[c])
+                    if cost_matrix[r, c] < (1 - self.match_threshold * 0.8):
+                        track_idx = unmatched_confirmed[r]
+                        det_idx = unmatched_dets[c]
+                        confirmed_tracks[track_idx].update(detections[det_idx])
+                        matched_det_indices.add(det_idx)
+                        
+                        # Remove from unmatched lists
+                        if track_idx in unmatched_confirmed:
+                            unmatched_confirmed.remove(track_idx)
+        
+        # Update unmatched detections list after confirmed track matching
+        unmatched_dets = [i for i in range(len(detections)) 
+                         if i not in matched_det_indices]
+        
+        # Third association: Tentative tracks
+        if unmatched_dets and tentative_tracks:
+            remaining_dets = [detections[i] for i in unmatched_dets]
+            
+            cost_matrix = self._calculate_enhanced_cost_matrix(
+                tentative_tracks, remaining_dets
+            )
+            
+            if cost_matrix.size > 0:
+                row_ind, col_ind = linear_sum_assignment(cost_matrix)
+                
+                for r, c in zip(row_ind, col_ind):
+                    if cost_matrix[r, c] < (1 - self.match_threshold * 0.6):
+                        det_idx = unmatched_dets[c]
+                        tentative_tracks[r].update(detections[det_idx])
+                        matched_det_indices.add(det_idx)
         
         # Mark unmatched tracks as missed
         for idx in unmatched_confirmed:
             confirmed_tracks[idx].mark_missed()
-        
+            
         for track in tentative_tracks:
             if track.time_since_update > 0:
                 track.mark_missed()
         
         # Create new tracks for unmatched detections
-        for det_idx in unmatched_dets:
-            # Check if detection is far from all existing tracks
+        final_unmatched_dets = [i for i in range(len(detections)) 
+                               if i not in matched_det_indices]
+        
+        for det_idx in final_unmatched_dets:
             detection = detections[det_idx]
-            is_new = True
             
-            # Additional check: ensure new detection is not too close to existing tracks
+            # Additional check: ensure not too close to existing tracks
+            is_new = True
             det_center = self._get_center(detection.bbox)
+            
             for track in self.tracks:
                 if track.state != 'deleted':
                     track_center = self._get_center(track.bbox)
-                    dist = np.sqrt((det_center[0] - track_center[0])**2 + 
-                                 (det_center[1] - track_center[1])**2)
+                    dist = np.linalg.norm(np.array(det_center) - np.array(track_center))
                     
-                    # If too close to existing track, likely same dog
-                    if dist < 50:  # Very close threshold
+                    if dist < 40:  # Very close threshold
                         is_new = False
                         break
             
@@ -202,44 +279,70 @@ class SimpleTracker:
         # Return only confirmed tracks
         return [t for t in self.tracks if t.state == 'confirmed']
     
-    def _calculate_cost_matrix(self, tracks: List[Track], 
-                               detections: List[Detection]) -> np.ndarray:
-        
+    def _calculate_enhanced_cost_matrix(self, tracks: List[Track],
+                                       detections: List[Detection]) -> np.ndarray:
+        """Calculate cost matrix with multiple cues"""
         if not tracks or not detections:
             return np.array([])
-        
-        matrix = np.ones((len(tracks), len(detections)))
+            
+        n_tracks = len(tracks)
+        n_dets = len(detections)
+        cost_matrix = np.ones((n_tracks, n_dets))
         
         for t_idx, track in enumerate(tracks):
-            track_center = self._get_center(track.bbox)
+            track_center = np.array(self._get_center(track.bbox))
+            track_size = np.array([
+                track.bbox[2] - track.bbox[0],
+                track.bbox[3] - track.bbox[1]
+            ])
             
             for d_idx, detection in enumerate(detections):
-                # Calculate IoU
+                # IoU cost
                 iou = self._iou(track.bbox, detection.bbox)
                 
-                # Calculate center distance
-                det_center = self._get_center(detection.bbox)
-                distance = np.sqrt((track_center[0] - det_center[0])**2 + 
-                                 (track_center[1] - det_center[1])**2)
+                # Center distance cost
+                det_center = np.array(self._get_center(detection.bbox))
+                distance = np.linalg.norm(track_center - det_center)
                 
-                # Combined cost (lower is better)
+                # Size similarity cost
+                det_size = np.array([
+                    detection.bbox[2] - detection.bbox[0],
+                    detection.bbox[3] - detection.bbox[1]
+                ])
+                size_sim = np.minimum(track_size, det_size) / (np.maximum(track_size, det_size) + 1e-6)
+                size_cost = 1 - np.mean(size_sim)
+                
+                # Combine costs
                 if iou >= self.min_iou_for_match and distance < self.max_center_distance:
-                    # Weighted combination: IoU is more important
                     iou_cost = 1 - iou
                     dist_cost = distance / self.max_center_distance
-                    matrix[t_idx, d_idx] = 0.7 * iou_cost + 0.3 * dist_cost
+                    
+                    # Weighted combination
+                    total_cost = (0.5 * iou_cost + 
+                                 0.3 * dist_cost + 
+                                 0.2 * size_cost)
+                    
+                    # Appearance cost if available
+                    if self.use_appearance and track.appearance_features and hasattr(detection, 'features'):
+                        # Simple cosine similarity
+                        track_feat = np.mean(track.appearance_features, axis=0)
+                        det_feat = detection.features
+                        app_sim = np.dot(track_feat, det_feat) / (
+                            np.linalg.norm(track_feat) * np.linalg.norm(det_feat) + 1e-6
+                        )
+                        app_cost = 1 - app_sim
+                        total_cost = 0.4 * iou_cost + 0.2 * dist_cost + 0.2 * size_cost + .2 * app_cost
+                    
+                    cost_matrix[t_idx, d_idx] = total_cost
                 else:
-                    # No match possible
-                    matrix[t_idx, d_idx] = 1.0
-                
-        return matrix
+                    cost_matrix[t_idx, d_idx] = 1.0
+                    
+        return cost_matrix
     
     def _get_center(self, bbox: List[float]) -> Tuple[float, float]:
-        
         return ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2)
     
     def _iou(self, bbox1: List[float], bbox2: List[float]) -> float:
-        
         x1 = max(bbox1[0], bbox2[0])
         y1 = max(bbox1[1], bbox2[1])
         x2 = min(bbox1[2], bbox2[2])
@@ -247,7 +350,7 @@ class SimpleTracker:
         
         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])
@@ -256,6 +359,10 @@ class SimpleTracker:
         return intersection / (union + 1e-6)
     
     def set_match_threshold(self, threshold: float):
-        
+        """Update matching threshold"""
         self.match_threshold = max(0.1, min(0.8, threshold))
-        
+        print(f"Tracking threshold updated to: {self.match_threshold:.2f}")
+
+
+# Use the enhanced tracker as default
+SimpleTracker = RobustTracker