Add back instances that were not tracked

talmolab · May 25, 2023 · 1ea047d · 1ea047d
1 parent d882640
commit 1ea047d
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/sleap/nn/tracking.py b/sleap/nn/tracking.py
@@ -1190,20 +1190,26 @@ def run_tracker(frames: List[LabeledFrame], tracker: BaseTracker) -> List[Labele
             inst.track = None
 
         # Prefer user instances over predicted instances
-        untracked_instances = (
-            lf.user_instances if lf.has_user_instances else lf.predicted_instances
-        )
-        track_args = {"untracked_instances": untracked_instances}
+        instances = []
+        if lf.has_user_instances:
+            instances_to_track = lf.user_instances
+            if lf.has_predicted_instances:
+                instances = lf.predicted_instances
+        else:
+            instances_to_track = lf.predicted_instances
+
+        track_args = {"untracked_instances": instances_to_track}
 
         if tracker.uses_image:
             track_args["img"] = lf.video[lf.frame_idx]
         else:
             track_args["img"] = None
 
+        instances.extend(tracker.track(**track_args))
         new_lf = LabeledFrame(
             frame_idx=lf.frame_idx,
             video=lf.video,
-            instances=tracker.track(**track_args),
+            instances=instances,
         )
         new_lfs.append(new_lf)