Update tracking.py

tracking.py

@@ -1,5 +1,5 @@
 """
-DeepSORT Tracker Implementation
+DeepSORT Tracker Implementation - Complete Fixed Version
 Uses Kalman filter for motion prediction + appearance features
 """
 import numpy as np
@@ -11,85 +11,53 @@ import uuid
 from detection import Detection
 
 class KalmanFilter:
-    """
-    Kalman Filter for tracking with constant velocity model
-    State: [x, y, w, h, vx, vy, vw, vh]
-    """
+    """Kalman Filter for tracking with constant velocity model"""
     def __init__(self):
-        # State transition matrix (constant velocity model)
         self.F = np.eye(8)
         for i in range(4):
             self.F[i, i+4] = 1
         
-        # Measurement matrix (we observe position and size)
         self.H = np.eye(4, 8)
-        
-        # Process noise
         self.Q = np.eye(8)
-        self.Q[4:, 4:] *= 0.01  # Smaller noise for velocities
-        
-        # Measurement noise
+        self.Q[4:, 4:] *= 0.01
         self.R = np.eye(4) * 10
-        
-        # State covariance
         self.P = np.eye(8) * 1000
-        
-        # State vector [x, y, w, h, vx, vy, vw, vh]
         self.x = np.zeros(8)
         
     def initiate(self, measurement: np.ndarray):
-        """Initialize filter with first measurement [x, y, w, h]"""
         self.x[:4] = measurement
-        self.x[4:] = 0  # Zero velocity
+        self.x[4:] = 0
         
     def predict(self):
-        """Predict next state"""
         self.x = self.F @ self.x
         self.P = self.F @ self.P @ self.F.T + self.Q
         return self.x[:4]
     
     def update(self, measurement: np.ndarray):
-        """Update with new measurement"""
-        # Innovation
         y = measurement - self.H @ self.x
-        
-        # Innovation covariance
         S = self.H @ self.P @ self.H.T + self.R
-        
-        # Kalman gain
         K = self.P @ self.H.T @ np.linalg.inv(S)
-        
-        # Update state
         self.x = self.x + K @ y
-        
-        # Update covariance
         self.P = (np.eye(8) - K @ self.H) @ self.P
-        
         return self.x[:4]
     
     def get_state(self) -> np.ndarray:
-        """Get current state [x, y, w, h]"""
         return self.x[:4]
 
 class DeepSORTTrack:
-    """
-    DeepSORT Track with Kalman filter and appearance features
-    """
+    """DeepSORT Track with Kalman filter and appearance features"""
     def __init__(self, detection: Detection, track_id: Optional[int] = None):
         self.track_id = track_id if track_id else self._generate_id()
         
-        # Convert bbox to center format [cx, cy, w, h]
         x1, y1, x2, y2 = detection.bbox
         cx = (x1 + x2) / 2
         cy = (y1 + y2) / 2
         w = x2 - x1
         h = y2 - y1
         
-        # Initialize Kalman filter
         self.kf = KalmanFilter()
         self.kf.initiate(np.array([cx, cy, w, h]))
         
-        # Track state
         self.detections = [detection]
         self.confidence = detection.confidence
         self.hits = 1
@@ -97,16 +65,13 @@ class DeepSORTTrack:
         self.time_since_update = 0
         self.state = 'tentative'
         
-        # Appearance features for ReID
-        self.features = deque(maxlen=100)  # Store last 100 features
+        self.features = deque(maxlen=100)
         if hasattr(detection, 'features') and detection.features is not None:
             self.features.append(detection.features)
         
-        # Trajectory
         self.trajectory = deque(maxlen=30)
         self.trajectory.append((cx, cy))
         
-        # Quality metrics
         self.avg_confidence = self.confidence
         self.consecutive_misses = 0
         
@@ -115,37 +80,24 @@ class DeepSORTTrack:
     
     @property
     def bbox(self) -> List[float]:
-        """Get bbox in [x1, y1, x2, y2] format"""
         cx, cy, w, h = self.kf.get_state()
-        return [
-            cx - w/2,
-            cy - h/2,
-            cx + w/2,
-            cy + h/2
-        ]
+        return [cx - w/2, cy - h/2, cx + w/2, cy + h/2]
     
     def predict(self):
-        """Predict next state using Kalman filter"""
         self.age += 1
         self.time_since_update += 1
         self.consecutive_misses += 1
-        
-        # Predict with Kalman filter
         self.kf.predict()
     
     def update(self, detection: Detection):
-        """Update track with new detection"""
-        # Convert to center format
         x1, y1, x2, y2 = detection.bbox
         cx = (x1 + x2) / 2
         cy = (y1 + y2) / 2
         w = x2 - x1
         h = y2 - y1
         
-        # Update Kalman filter
         self.kf.update(np.array([cx, cy, w, h]))
         
-        # Update track state
         self.detections.append(detection)
         self.confidence = detection.confidence
         self.avg_confidence = 0.9 * self.avg_confidence + 0.1 * self.confidence
@@ -154,18 +106,14 @@ class DeepSORTTrack:
         self.time_since_update = 0
         self.consecutive_misses = 0
         
-        # Update features
         if hasattr(detection, 'features') and detection.features is not None:
             self.features.append(detection.features)
         
-        # Update trajectory
         self.trajectory.append((cx, cy))
         
-        # Confirm track
         if self.state == 'tentative' and self.hits >= 3:
             self.state = 'confirmed'
         
-        # Keep only recent detections
         if len(self.detections) > 5:
             for old_det in self.detections[:-5]:
                 if hasattr(old_det, 'image_crop'):
@@ -173,7 +121,6 @@ class DeepSORTTrack:
             self.detections = self.detections[-5:]
     
     def mark_missed(self):
-        """Mark track as missed"""
         if self.state == 'confirmed':
             if self.consecutive_misses > 30 or self.time_since_update > 60:
                 self.state = 'deleted'
@@ -182,34 +129,21 @@ class DeepSORTTrack:
                 self.state = 'deleted'
     
     def get_feature(self) -> Optional[np.ndarray]:
-        """Get averaged appearance feature"""
         if not self.features:
             return None
-        # Use exponential moving average for recent features
         features_array = np.array(list(self.features))
         weights = np.exp(np.linspace(-1, 0, len(features_array)))
         weights /= weights.sum()
         return np.average(features_array, axis=0, weights=weights)
 
 class DeepSORTTracker:
-    """
-    DeepSORT Tracker combining Kalman filter motion model
-    with appearance feature matching
-    """
+    """DeepSORT Tracker combining Kalman filter with appearance features"""
     def __init__(self,
                  max_iou_distance: float = 0.7,
                  max_age: int = 30,
                  n_init: int = 3,
                  nn_budget: int = 100,
                  use_appearance: bool = True):
-        """
-        Args:
-            max_iou_distance: Maximum IoU distance for matching (0.7 = 30% IoU)
-            max_age: Maximum frames to keep lost tracks
-            n_init: Number of frames to confirm a track
-            nn_budget: Maximum size of appearance feature gallery
-            use_appearance: Whether to use appearance features
-        """
         self.max_iou_distance = max_iou_distance
         self.max_age = max_age
         self.n_init = n_init
@@ -219,65 +153,46 @@ class DeepSORTTracker:
         self.tracks: List[DeepSORTTrack] = []
         self.track_id_count = 1
         
-        # Gating thresholds
-        self.gating_threshold_iou = 0.3  # Minimum IoU for association
-        self.gating_threshold_appearance = 0.5  # Maximum cosine distance
+        self.gating_threshold_iou = 0.3
+        self.gating_threshold_appearance = 0.5
         
     def update(self, detections: List[Detection]) -> List[DeepSORTTrack]:
-        """Update tracks with new detections"""
-        # Predict existing tracks
         for track in self.tracks:
             track.predict()
         
-        # Match detections to tracks
         matches, unmatched_tracks, unmatched_detections = self._match(
             detections, self.tracks
         )
         
-        # Update matched tracks
         for track_idx, det_idx in matches:
             self.tracks[track_idx].update(detections[det_idx])
         
-        # Mark unmatched tracks as missed
         for track_idx in unmatched_tracks:
             self.tracks[track_idx].mark_missed()
         
-        # Create new tracks for unmatched detections
         for det_idx in unmatched_detections:
             self._initiate_track(detections[det_idx])
         
-        # 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 _match(self, detections: List[Detection], tracks: List[DeepSORTTrack]) -> Tuple:
-        """
-        Match detections to tracks using cascade matching
-        Returns: (matches, unmatched_tracks, unmatched_detections)
-        """
         if not tracks or not detections:
             return [], list(range(len(tracks))), list(range(len(detections)))
         
-        # Split tracks by state
         confirmed_tracks = [i for i, t in enumerate(tracks) if t.state == 'confirmed']
         unconfirmed_tracks = [i for i, t in enumerate(tracks) if t.state == 'tentative']
         
-        # Stage 1: Match confirmed tracks
         matches_a, unmatched_tracks_a, unmatched_detections = \
             self._matching_cascade(detections, tracks, confirmed_tracks)
         
-        # Stage 2: Match unconfirmed tracks
         iou_track_candidates = unconfirmed_tracks + unmatched_tracks_a
-        
-        # Filter detections used in stage 1
         remaining_detections = [detections[i] for i in unmatched_detections]
         
         matches_b, unmatched_tracks_b, unmatched_detections_b = \
             self._match_iou(remaining_detections, tracks, iou_track_candidates)
         
-        # Remap indices for stage 2
         matches_b = [(t, unmatched_detections[d]) for t, d in matches_b]
         unmatched_detections = [unmatched_detections[i] for i in unmatched_detections_b]
         
@@ -289,11 +204,9 @@ class DeepSORTTracker:
     def _matching_cascade(self, detections: List[Detection], 
                           tracks: List[DeepSORTTrack],
                           track_indices: List[int]) -> Tuple:
-        """Cascade matching by time since update"""
         matches = []
         unmatched_detections = list(range(len(detections)))
         
-        # Group tracks by time since update
         for level in range(self.max_age):
             if len(unmatched_detections) == 0:
                 break
@@ -306,22 +219,18 @@ class DeepSORTTracker:
             if len(track_indices_l) == 0:
                 continue
             
-            # Match at this level
-            matches_l, _, unmatched_detections = self._match_features_and_iou(
-                [detections[i] for i in unmatched_detections],
+            detection_subset = [detections[i] for i in unmatched_detections]
+            
+            matches_l, _, unmatched_subset_indices = self._match_features_and_iou(
+                detection_subset,
                 tracks,
                 track_indices_l
             )
             
-            # Remap detection indices
-            matches_l = [(t, unmatched_detections[d]) for t, d in matches_l]
-            matches.extend(matches_l)
+            matches_l_remapped = [(t, unmatched_detections[d]) for t, d in matches_l]
+            matches.extend(matches_l_remapped)
             
-            # Update unmatched detections
-            unmatched_detections = [
-                unmatched_detections[i] for i in range(len(unmatched_detections))
-                if i not in [d for _, d in matches_l]
-            ]
+            unmatched_detections = [unmatched_detections[i] for i in unmatched_subset_indices]
         
         unmatched_tracks = [
             k for k in track_indices
@@ -333,11 +242,9 @@ class DeepSORTTracker:
     def _match_features_and_iou(self, detections: List[Detection],
                                  tracks: List[DeepSORTTrack],
                                  track_indices: List[int]) -> Tuple:
-        """Match using both appearance features and IoU"""
         if not track_indices or not detections:
             return [], track_indices, list(range(len(detections)))
         
-        # Calculate cost matrix
         cost_matrix = np.zeros((len(track_indices), len(detections)))
         
         for i, track_idx in enumerate(track_indices):
@@ -348,27 +255,23 @@ class DeepSORTTracker:
             for j, detection in enumerate(detections):
                 det_bbox = detection.bbox
                 
-                # IoU cost
                 iou = self._iou(track_bbox, det_bbox)
                 iou_cost = 1 - iou
                 
-                # Appearance cost
                 if self.use_appearance and track_feature is not None:
                     if hasattr(detection, 'features') and detection.features is not None:
                         cosine_dist = cosine(track_feature, detection.features)
                         appearance_cost = cosine_dist
                     else:
-                        appearance_cost = 0.5  # Neutral if no features
+                        appearance_cost = 0.5
                 else:
                     appearance_cost = 0.0
                 
-                # Combined cost (weighted)
                 if self.use_appearance and track_feature is not None:
                     cost = 0.5 * iou_cost + 0.5 * appearance_cost
                 else:
                     cost = iou_cost
                 
-                # Gating: set to infinity if outside thresholds
                 if iou < self.gating_threshold_iou:
                     cost = 1e6
                 if self.use_appearance and appearance_cost > self.gating_threshold_appearance:
@@ -376,7 +279,6 @@ class DeepSORTTracker:
                 
                 cost_matrix[i, j] = cost
         
-        # Hungarian algorithm
         row_indices, col_indices = linear_sum_assignment(cost_matrix)
         
         matches = []
@@ -396,11 +298,9 @@ class DeepSORTTracker:
     def _match_iou(self, detections: List[Detection],
                    tracks: List[DeepSORTTrack],
                    track_indices: List[int]) -> Tuple:
-        """Simple IoU matching for unconfirmed tracks"""
         if not track_indices or not detections:
             return [], track_indices, list(range(len(detections)))
         
-        # IoU distance matrix
         iou_matrix = np.zeros((len(track_indices), len(detections)))
         
         for i, track_idx in enumerate(track_indices):
@@ -409,7 +309,6 @@ class DeepSORTTracker:
                 iou = self._iou(track.bbox, detection.bbox)
                 iou_matrix[i, j] = 1 - iou
         
-        # Hungarian matching
         row_indices, col_indices = linear_sum_assignment(iou_matrix)
         
         matches = []
@@ -427,13 +326,11 @@ class DeepSORTTracker:
         return matches, unmatched_tracks, unmatched_detections
     
     def _initiate_track(self, detection: Detection):
-        """Create new track"""
         new_track = DeepSORTTrack(detection, self.track_id_count)
         self.track_id_count += 1
         self.tracks.append(new_track)
     
     def _iou(self, bbox1: List[float], bbox2: List[float]) -> float:
-        """Calculate IoU between two bboxes"""
         try:
             x1 = max(bbox1[0], bbox2[0])
             y1 = max(bbox1[1], bbox2[1])
@@ -453,11 +350,9 @@ class DeepSORTTracker:
             return 0.0
     
     def reset(self):
-        """Reset tracker"""
         self.tracks.clear()
         self.track_id_count = 1
         print("DeepSORT tracker reset")
 
-# Compatibility alias
 SimpleTracker = DeepSORTTracker
 RobustTracker = DeepSORTTracker