fix small edge case when initializing new tracks

talmolab · May 28, 2024 · e1f7c66 · e1f7c66
1 parent 9d4d24a
commit e1f7c66
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/biogtr/inference/tracker.py b/biogtr/inference/tracker.py
@@ -156,7 +156,7 @@ def sliding_inference(self, model: GlobalTrackingTransformer, frames: list[Frame
                 self.track_queue.end_tracks()
 
             """
-            Initialize tracks on first frame of video or first instance of detections.
+            Initialize tracks on first frame where detections appear.
             """
             if len(self.track_queue) == 0:
                 if frame_to_track.has_instances():
@@ -168,12 +168,13 @@ def sliding_inference(self, model: GlobalTrackingTransformer, frames: list[Frame
                     curr_track_id = 0
                     for i, instance in enumerate(frames[batch_idx].instances):
                         instance.pred_track_id = instance.gt_track_id
-                        curr_track_id = instance.pred_track_id
+                        curr_track_id = max(curr_track_id, instance.pred_track_id)
 
                     for i, instance in enumerate(frames[batch_idx].instances):
                         if instance.pred_track_id == -1:
-                            instance.pred_track_id = curr_track_id
                             curr_track += 1
+                            instance.pred_track_id = curr_track_id
+
 
             else:
                 if (
@@ -251,6 +252,7 @@ def _run_global_tracker(
         overlap_thresh = self.overlap_thresh
         mult_thresh = self.mult_thresh
         n_traj = self.track_queue.n_tracks
+        curr_track = self.track_queue.curr_track
 
         reid_features = torch.cat([frame.get_features() for frame in frames], dim=0)[
             None
@@ -471,8 +473,9 @@ def _run_global_tracker(
             if track_ids[i] < 0:
                 if self.verbose:
                     print(f"Creating new track {n_traj}")
-                track_ids[i] = n_traj
-                n_traj += 1
+                curr_track += 1
+                track_ids[i] = curr_track
+
 
         query_frame.matches = (match_i, match_j)