Fix target instance count parameter (#358)

Currently, enabling `post_connect_single_breaks` forces users to set
`tracking_target_instance_count`, even when `max_instances` is already
provided. In this PR we fix this by ensuring when
post_connect_single_breaks=True and tracking_target_instance_count is
None, we infer the target from max_instances; if
tracking_target_instance_count is explicitly set, it still takes
precedence.

Author: gitttt-1234

pyproject.toml

@@ -54,19 +54,19 @@ readme = {file = ["README.md"], content-type="text/markdown"}
 [project.optional-dependencies]
 torch = [
     "torch",
-    "torchvision",
+    "torchvision<0.24.0",
 ]
 torch-cpu = [
     "torch",
-    "torchvision",
+    "torchvision<0.24.0",
 ]
 torch-cuda118 = [
     "torch",
-    "torchvision",
+    "torchvision<0.24.0",
 ]
 torch-cuda128 = [
     "torch",
-    "torchvision",
+    "torchvision<0.24.0",
 ]
 dev = [
     "pytest",

sleap_nn/cli.py

@@ -389,7 +389,7 @@ def train(config_name, config_dir, overrides):
 @click.option(
     "--tracking_target_instance_count",
     type=int,
-    default=0,
+    default=None,
     help="Target number of instances to track per frame. (default: 0)",
 )
 @click.option(

sleap_nn/predict.py

@@ -98,7 +98,7 @@ def run_inference(
     of_window_size: int = 21,
     of_max_levels: int = 3,
     post_connect_single_breaks: bool = False,
-    tracking_target_instance_count: int = 0,
+    tracking_target_instance_count: Optional[int] = None,
     tracking_pre_cull_to_target: int = 0,
     tracking_pre_cull_iou_threshold: float = 0,
     tracking_clean_instance_count: int = 0,
@@ -236,7 +236,7 @@ def run_inference(
             Default: 3. (only if `use_flow` is True).
         post_connect_single_breaks: If True and `max_tracks` is not None with local queues candidate method,
             connects track breaks when exactly one track is lost and exactly one new track is spawned in the frame.
-        tracking_target_instance_count: Target number of instances to track per frame. (default: 0)
+        tracking_target_instance_count: Target number of instances to track per frame. (default: None)
         tracking_pre_cull_to_target: If non-zero and target_instance_count is also non-zero, then cull instances over target count per frame *before* tracking. (default: 0)
         tracking_pre_cull_iou_threshold: If non-zero and pre_cull_to_target also set, then use IOU threshold to remove overlapping instances over count *before* tracking. (default: 0)
         tracking_clean_instance_count: Target number of instances to clean *after* tracking. (default: 0)
@@ -300,6 +300,14 @@ def run_inference(
 
             logger.info(f"Running tracking on {len(lf_frames)} frames...")
 
+            if post_connect_single_breaks or tracking_pre_cull_to_target:
+                if tracking_target_instance_count is None and max_instances is None:
+                    message = "Both tracking_target_instance_count and max_instances is set to 0. To connect single breaks or pre-cull to target, at least one of them should be set to an integer."
+                    logger.error(message)
+                    raise ValueError(message)
+                elif tracking_target_instance_count is None:
+                    tracking_target_instance_count = max_instances
+
             tracked_frames = run_tracker(
                 untracked_frames=lf_frames,
                 window_size=tracking_window_size,
@@ -377,6 +385,13 @@ def run_inference(
             and not isinstance(predictor, BottomUpMultiClassPredictor)
             and not isinstance(predictor, TopDownMultiClassPredictor)
         ):
+            if post_connect_single_breaks or tracking_pre_cull_to_target:
+                if tracking_target_instance_count is None and max_instances is None:
+                    message = "Both tracking_target_instance_count and max_instances is set to 0. To connect single breaks or pre-cull to target, at least one of them should be set to an integer."
+                    logger.error(message)
+                    raise ValueError(message)
+                elif tracking_target_instance_count is None:
+                    tracking_target_instance_count = max_instances
             predictor.tracker = Tracker.from_config(
                 candidates_method=candidates_method,
                 min_match_points=min_match_points,
@@ -454,11 +469,6 @@ def run_inference(
                         lfs, tracking_clean_instance_count
                     )
             elif post_connect_single_breaks:
-                if not tracking_target_instance_count:
-                    message = "tracking_target_instance_count is 0. To connect single breaks, tracking_target_instance_count should be set to an integer."
-                    logger.error(message)
-                    raise ValueError(message)
-
                 start_final_pass_time = time()
                 start_fp_timestamp = str(datetime.now())
                 logger.info(

sleap_nn/tracking/tracker.py

@@ -76,7 +76,7 @@ class Tracker:
     robust_best_instance: float = 1.0
     use_flow: bool = False
     is_local_queue: bool = False
-    tracking_target_instance_count: int = 0
+    tracking_target_instance_count: Optional[int] = None
     tracking_pre_cull_to_target: int = 0
     tracking_pre_cull_iou_threshold: float = 0
     _scoring_functions: Dict[str, Any] = {
@@ -119,7 +119,7 @@ def from_config(
         of_img_scale: float = 1.0,
         of_window_size: int = 21,
         of_max_levels: int = 3,
-        tracking_target_instance_count: int = 0,
+        tracking_target_instance_count: Optional[int] = None,
         tracking_pre_cull_to_target: int = 0,
         tracking_pre_cull_iou_threshold: float = 0,
     ):
@@ -162,7 +162,7 @@ def from_config(
             of_max_levels: Number of pyramid scale levels to consider. This is different
                 from the scale parameter, which determines the initial image scaling.
                 Default: 3. (only if `use_flow` is True)
-            tracking_target_instance_count: Target number of instances to track per frame. (default: 0)
+            tracking_target_instance_count: Target number of instances to track per frame. (default: None)
             tracking_pre_cull_to_target: If non-zero and target_instance_count is also non-zero, then cull instances over target count per frame *before* tracking. (default: 0)
             tracking_pre_cull_iou_threshold: If non-zero and pre_cull_to_target also set, then use IOU threshold to remove overlapping instances over count *before* tracking. (default: 0)
 
@@ -237,7 +237,11 @@ def track(
         Returns:
             List of `sio.PredictedInstance` objects, each having an assigned track.
         """
-        if self.tracking_target_instance_count and self.tracking_pre_cull_to_target:
+        if (
+            self.tracking_target_instance_count is not None
+            and self.tracking_target_instance_count
+            and self.tracking_pre_cull_to_target
+        ):
             untracked_instances = cull_frame_instances(
                 untracked_instances,
                 self.tracking_target_instance_count,
@@ -492,7 +496,7 @@ class FlowShiftTracker(Tracker):
         of_max_levels: Number of pyramid scale levels to consider. This is different
             from the scale parameter, which determines the initial image scaling.
             Default: 3
-        tracking_target_instance_count: Target number of instances to track per frame. (default: 0)
+        tracking_target_instance_count: Target number of instances to track per frame. (default: None)
         tracking_pre_cull_to_target: If non-zero and target_instance_count is also non-zero, then cull instances over target count per frame *before* tracking. (default: 0)
         tracking_pre_cull_iou_threshold: If non-zero and pre_cull_to_target also set, then use IOU threshold to remove overlapping instances over count *before* tracking. (default: 0)
 
@@ -729,7 +733,7 @@ def run_tracker(
     of_window_size: int = 21,
     of_max_levels: int = 3,
     post_connect_single_breaks: bool = False,
-    tracking_target_instance_count: int = 0,
+    tracking_target_instance_count: Optional[int] = None,
     tracking_pre_cull_to_target: int = 0,
     tracking_pre_cull_iou_threshold: float = 0,
     tracking_clean_instance_count: int = 0,
@@ -777,7 +781,7 @@ def run_tracker(
                 Default: 3. (only if `use_flow` is True).
         post_connect_single_breaks: If True and `max_tracks` is not None with local queues candidate method,
             connects track breaks when exactly one track is lost and exactly one new track is spawned in the frame.
-        tracking_target_instance_count: Target number of instances to track per frame. (default: 0)
+        tracking_target_instance_count: Target number of instances to track per frame. (default: None)
         tracking_pre_cull_to_target: If non-zero and target_instance_count is also non-zero, then cull instances over target count per frame *before* tracking. (default: 0)
         tracking_pre_cull_iou_threshold: If non-zero and pre_cull_to_target also set, then use IOU threshold to remove overlapping instances over count *before* tracking. (default: 0)
         tracking_clean_instance_count: Target number of instances to clean *after* tracking. (default: 0)
@@ -840,8 +844,11 @@ def run_tracker(
             )
 
     if post_connect_single_breaks:
-        if not tracking_target_instance_count:
-            message = "tracking_target_instance_count is 0. To connect single breaks, tracking_target_instance_count should be set to an integer."
+        if (
+            tracking_target_instance_count is None
+            or tracking_target_instance_count == 0
+        ):
+            message = "tracking_target_instance_count is None or 0. To connect single breaks, tracking_target_instance_count should be set to an integer."
             logger.error(message)
             raise ValueError(message)
         start_final_pass_time = time()

tests/test_predict.py

@@ -260,6 +260,31 @@ def test_topdown_predictor(
 
     assert len(pred_labels.tracks) <= 2  # should be less than max tracks
 
+    for lf in pred_labels:
+        for instance in lf.instances:
+            assert instance.track is not None
+
+    # test with tracking and tracking_target_instance_count is not provided
+    pred_labels = run_inference(
+        model_paths=[
+            minimal_instance_centroid_ckpt,
+            minimal_instance_centered_instance_ckpt,
+        ],
+        data_path=centered_instance_video.as_posix(),
+        make_labels=True,
+        output_path=tmp_path,
+        max_instances=2,
+        post_connect_single_breaks=True,
+        max_tracks=None,
+        device="cpu",
+        peak_threshold=0.1,
+        frames=[x for x in range(5)],
+        tracking=True,
+        integral_refinement=None,
+    )
+
+    assert len(pred_labels.tracks) <= 2  # should be less than max tracks
+
     for lf in pred_labels:
         for instance in lf.instances:
             assert instance.track is not None
@@ -276,15 +301,14 @@ def test_topdown_predictor(
             output_path=tmp_path,
             max_instances=None,
             post_connect_single_breaks=True,
-            tracking_target_instance_count=0,
             max_tracks=None,
             device="cpu",
             peak_threshold=0.1,
             frames=[x for x in range(20)],
             tracking=True,
             integral_refinement=None,
         )
-        assert "tracking_target_instance_count is 0" in caplog.text
+        assert "tracking_target_instance_count and max_instances" in caplog.text
 
 
 def test_multiclass_topdown_predictor(
@@ -1050,6 +1074,27 @@ def test_tracking_only_pipeline(
             integral_refinement=None,
         )
 
+    # neither max_instances nor tracking_target_instance_count is provided
+    with pytest.raises(ValueError):
+        labels = run_inference(
+            data_path=centered_instance_video.as_posix(),
+            tracking=True,
+            integral_refinement=None,
+            post_connect_single_breaks=True,
+        )
+
+    # racking_target_instance_count is provided
+    labels = run_inference(
+        data_path=minimal_instance.as_posix(),
+        tracking=True,
+        integral_refinement=None,
+        post_connect_single_breaks=True,
+        max_instances=2,
+        output_path=tmp_path,
+    )
+    for lf in labels:
+        assert len(lf.instances) <= 2
+
 
 def test_legacy_topdown_predictor(
     minimal_instance,