Refactor ByteTracker: add min_hits confirmation, split thresholds, associate unconfirmed tracks

utils/tracker.py

@@ -223,17 +223,18 @@ class STrack:
         # wait, input is xyxy usually in our pipeline
         # ByteTrack usually uses tlwh internally.
         # Let's standardize to input xyxy.
-        
+
         self._tlwh = np.asarray(self._tlwh_from_xyxy(tlwh), dtype=np.float32)
         self.is_activated = False
         self.track_id = 0
         self.state = 1 # 1: New, 2: Tracked, 3: Lost, 4: Removed
-        
+
         self.score = score
         self.label = label
         self.start_frame = 0
         self.frame_id = 0
         self.time_since_update = 0
+        self.hits = 0
         
         # Multi-frame history
         self.history = []
@@ -282,15 +283,16 @@ class STrack:
         return ret
 
     def activate(self, kalman_filter, frame_id):
-        """Start a new track."""
+        """Start a new track (tentative until min_hits reached)."""
         self.kalman_filter = kalman_filter
         self.track_id = self.next_id()
         self.mean, self.covariance = self.kalman_filter.initiate(self.tlbr) # Initiate needs xyxy
-        
+
         self.state = 2 # Tracked
         self.frame_id = frame_id
         self.start_frame = frame_id
-        self.is_activated = True
+        self.is_activated = False  # Tentative until min_hits reached
+        self.hits = 1
 
     def re_activate(self, new_track, frame_id, new_id=False):
         """Reactivate a lost track with a new detection."""
@@ -305,17 +307,19 @@ class STrack:
         if new_id:
             self.track_id = self.next_id()
 
-    def update(self, new_track, frame_id):
+    def update(self, new_track, frame_id, min_hits=3):
         """Update a tracked object with a new detection."""
         self.frame_id = frame_id
         self.time_since_update = 0
         self.score = new_track.score
-        
+        self.hits += 1
+
         self.mean, self.covariance = self.kalman_filter.update(
             self.mean, self.covariance, self._xyah_from_xyxy(new_track.tlbr)
         )
         self.state = 2 # Tracked
-        self.is_activated = True
+        if self.hits >= min_hits:
+            self.is_activated = True  # Confirmed after N consecutive hits
         
     def predict(self):
         """Propagate tracking state distribution one time step forward."""
@@ -355,13 +359,18 @@ class STrack:
 
 
 class ByteTracker:
-    def __init__(self, track_thresh=0.3, track_buffer=60, match_thresh=0.8, frame_rate=30):
+    def __init__(self, track_high_thresh=0.3, track_low_thresh=0.1,
+                 new_track_thresh=0.4, track_buffer=60, match_thresh=0.8,
+                 frame_rate=30, min_hits=3):
         STrack.reset_count()
-        self.track_thresh = track_thresh
+        self.track_high_thresh = track_high_thresh
+        self.track_low_thresh = track_low_thresh
+        self.new_track_thresh = new_track_thresh
         self.track_buffer = track_buffer
         self.match_thresh = match_thresh
+        self.min_hits = min_hits
         self.frame_id = 0
-        
+
         self.tracked_stracks = []  # Type: List[STrack]
         self.lost_stracks = []     # Type: List[STrack]
         self.removed_stracks = []  # Type: List[STrack]
@@ -401,14 +410,14 @@ class ByteTracker:
         # We wrap original dict in STrack
         
         for d in detections_list:
-            bbox = d['bbox']
             score = d['score']
-            label = d['label']
-            
-            t = STrack(bbox, score, label)
-            t.original_data = d # Link back
-            
-            if score >= self.track_thresh:
+            if score < self.track_low_thresh:
+                continue  # Background noise — discard entirely
+
+            t = STrack(d['bbox'], score, d['label'])
+            t.original_data = d  # Link back
+
+            if score >= self.track_high_thresh:
                 detections.append(t)
             else:
                 detections_second.append(t)
@@ -435,12 +444,12 @@ class ByteTracker:
             track = strack_pool[itracked]
             det = detections[idet]
             if track.state == 2:
-                track.update(det, self.frame_id)
+                track.update(det, self.frame_id, min_hits=self.min_hits)
                 activated_stracks.append(track)
             else:
                 track.re_activate(det, self.frame_id, new_id=False)
                 refind_stracks.append(track)
-            
+
             # Persist data
             self._sync_data(track, det)
 
@@ -454,12 +463,12 @@ class ByteTracker:
             track = r_tracked_stracks[itracked]
             det = detections_second[idet]
             if track.state == 2:
-                track.update(det, self.frame_id)
+                track.update(det, self.frame_id, min_hits=self.min_hits)
                 activated_stracks.append(track)
             else:
                 track.re_activate(det, self.frame_id, new_id=False)
                 refind_stracks.append(track)
-            
+
             self._sync_data(track, det)
 
         for it in u_track:
@@ -468,12 +477,41 @@ class ByteTracker:
                 track.state = 3 # Lost
                 lost_stracks.append(track)
 
+        # 3.5 Associate unconfirmed tracks with remaining unmatched detections
+        if unconfirmed and u_detection:
+            remaining_dets = [detections[i] for i in u_detection]
+            dists = iou_distance(unconfirmed, remaining_dets)
+            matches_unc, u_unconfirmed, u_det_remaining = linear_assignment(dists, thresh=0.7)
+
+            for itracked, idet in matches_unc:
+                track = unconfirmed[itracked]
+                det = remaining_dets[idet]
+                track.update(det, self.frame_id, min_hits=self.min_hits)
+                activated_stracks.append(track)
+                self._sync_data(track, det)
+
+            # Update u_detection to only contain indices not matched to unconfirmed
+            matched_det_indices = set(u_detection[idet] for _, idet in matches_unc) if len(matches_unc) > 0 else set()
+            u_detection = [i for i in u_detection if i not in matched_det_indices]
+
+            # Unconfirmed tracks that didn't match → remove (too noisy to keep)
+            for it in u_unconfirmed:
+                track = unconfirmed[it]
+                track.state = 4  # Removed
+                removed_stracks.append(track)
+
+        elif unconfirmed:
+            # No detections left to match — remove all unconfirmed
+            for track in unconfirmed:
+                track.state = 4  # Removed
+                removed_stracks.append(track)
+
         # 4. Init new tracks from unmatched high score detections
         # Note: Unmatched low score detections are ignored (noise)
         unmatched_dets = [detections[i] for i in u_detection]
         for track in unmatched_dets:
-            if track.score < self.track_thresh:
-                continue
+            if track.score < self.new_track_thresh:
+                continue  # Not confident enough to start a new track
 
             track.activate(self.kalman_filter, self.frame_id)
             activated_stracks.append(track)