Add track cleaning args (#349)

This PR adds track cleaning arguments (pre culling instances, and post
tracking clean-up arguments).

Author: gitttt-1234

docs/inference.md

@@ -107,6 +107,8 @@ sleap-nn track \
 | `--video_dataset` | The dataset for HDF5 videos | `None` |
 | `--video_input_format` | The input_format for HDF5 videos | `channels_last` |
 | `--frames` | List of frames indices. If `None`, all frames in the video are used | All frames |
+| `--no_empty_frames` | If `True`, removes frames with no predicted instances from the output labels | `False` |
+
 
 #### Performance
 
@@ -177,6 +179,14 @@ When using the `sleap-nn track` CLI command with both `--model_paths` and `--tra
 | `--of_max_levels` | Number of pyramid scale levels to consider. This is different from the scale parameter, which determines the initial image scaling (only if `use_flow` is True) | `3` |
 | `--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 | `False` |
 
+!!! warning "Tracking cleaning and pre-cull parameters"
+
+    The parameters `--tracking_pre_cull_to_target`, `--tracking_target_instance_count`, `tracking_pre_cull_iou_threshold`, `tracking_clean_iou_threshold` and `--tracking_clean_instance_count` are provided for backwards compatibility with legacy SLEAP workflows and **may be deprecated in future releases**. 
+
+    - To restrict the number of instances per frame, use the `--max_instances` parameter, which selects the top instances with the highest prediction scores.
+
+    We recommend using `--max_instances` for controlling the number of predicted instances per frame in new projects.
+
 #### Fixed Window Tracking
 
 This method maintains a fixed-size queue with the last N frames and uses all instances from those frames as candidates for matching.

sleap_nn/cli.py

@@ -209,6 +209,12 @@ def train(config_name, config_dir, overrides):
     default=False,
     help="Only run inference on unlabeled suggested frames when running on labels dataset. This is useful for generating predictions for initialization during labeling.",
 )
+@click.option(
+    "--no_empty_frames",
+    is_flag=True,
+    default=False,
+    help=("Clear empty frames that did not have predictions before saving to output."),
+)
 @click.option(
     "--video_index",
     type=int,
@@ -380,6 +386,42 @@ def train(config_name, config_dir, overrides):
     default=False,
     help="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.",
 )
+@click.option(
+    "--tracking_target_instance_count",
+    type=int,
+    default=0,
+    help="Target number of instances to track per frame. (default: 0)",
+)
+@click.option(
+    "--tracking_pre_cull_to_target",
+    type=int,
+    default=0,
+    help=(
+        "If non-zero and target_instance_count is also non-zero, then cull instances "
+        "over target count per frame *before* tracking. (default: 0)"
+    ),
+)
+@click.option(
+    "--tracking_pre_cull_iou_threshold",
+    type=float,
+    default=0,
+    help=(
+        "If non-zero and pre_cull_to_target also set, then use IOU threshold to remove "
+        "overlapping instances over count *before* tracking. (default: 0)"
+    ),
+)
+@click.option(
+    "--tracking_clean_instance_count",
+    type=int,
+    default=0,
+    help="Target number of instances to clean *after* tracking. (default: 0)",
+)
+@click.option(
+    "--tracking_clean_iou_threshold",
+    type=float,
+    default=0,
+    help="IOU to use when culling instances *after* tracking. (default: 0)",
+)
 def track(**kwargs):
     """Run Inference and Tracking workflow."""
     # Convert model_paths from tuple to list

sleap_nn/predict.py

@@ -9,7 +9,12 @@
     BottomUpMultiClassPredictor,
     TopDownMultiClassPredictor,
 )
-from sleap_nn.tracking.tracker import Tracker, run_tracker, connect_single_breaks
+from sleap_nn.tracking.tracker import (
+    Tracker,
+    run_tracker,
+    connect_single_breaks,
+    cull_instances,
+)
 from omegaconf import OmegaConf
 import sleap_io as sio
 from pathlib import Path
@@ -53,6 +58,7 @@ def run_inference(
     anchor_part: Optional[str] = None,
     only_labeled_frames: bool = False,
     only_suggested_frames: bool = False,
+    no_empty_frames: bool = False,
     batch_size: int = 4,
     queue_maxsize: int = 8,
     video_index: Optional[int] = None,
@@ -92,6 +98,11 @@ 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_pre_cull_to_target: int = 0,
+    tracking_pre_cull_iou_threshold: float = 0,
+    tracking_clean_instance_count: int = 0,
+    tracking_clean_iou_threshold: float = 0,
 ):
     """Entry point to run inference on trained SLEAP-NN models.
 
@@ -125,6 +136,7 @@ def run_inference(
                 provided, the anchor part in the `training_config.yaml` is used. Default: `None`.
         only_labeled_frames: (bool) `True` if inference should be run only on user-labeled frames. Default: `False`.
         only_suggested_frames: (bool) `True` if inference should be run only on unlabeled suggested frames. Default: `False`.
+        no_empty_frames: (bool) `True` if empty frames that did not have predictions should be cleared before saving to output. Default: `False`.
         batch_size: (int) Number of samples per batch. Default: 4.
         queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 8.
         video_index: (int) Integer index of video in .slp file to predict on. To be used with
@@ -224,6 +236,11 @@ 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_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)
+        tracking_clean_iou_threshold: IOU to use when culling instances *after* tracking. (default: 0)
 
     Returns:
         Returns `sio.Labels` object if `make_labels` is True. Else this function returns
@@ -300,6 +317,11 @@ def run_inference(
                 of_window_size=of_window_size,
                 of_max_levels=of_max_levels,
                 post_connect_single_breaks=post_connect_single_breaks,
+                tracking_target_instance_count=tracking_target_instance_count,
+                tracking_pre_cull_to_target=tracking_pre_cull_to_target,
+                tracking_pre_cull_iou_threshold=tracking_pre_cull_iou_threshold,
+                tracking_clean_instance_count=tracking_clean_instance_count,
+                tracking_clean_iou_threshold=tracking_clean_iou_threshold,
             )
 
             finish_timestamp = str(datetime.now())
@@ -370,6 +392,9 @@ def run_inference(
                 of_img_scale=of_img_scale,
                 of_window_size=of_window_size,
                 of_max_levels=of_max_levels,
+                tracking_target_instance_count=tracking_target_instance_count,
+                tracking_pre_cull_to_target=tracking_pre_cull_to_target,
+                tracking_pre_cull_iou_threshold=tracking_pre_cull_iou_threshold,
             )
 
         if isinstance(predictor, BottomUpPredictor):
@@ -418,29 +443,38 @@ def run_inference(
             make_labels=make_labels,
         )
 
-        if tracking and post_connect_single_breaks:
-            if max_tracks is None:
-                max_tracks = max_instances
-
-            if max_tracks is None:
-                message = "Max_tracks (and max instances) is None. To connect single breaks, max_tracks 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(
-                f"Started final-pass (connecting single breaks) at: {start_fp_timestamp}"
-            )
-            corrected_lfs = connect_single_breaks(
-                lfs=[x for x in output], max_instances=max_tracks
-            )
-            finish_fp_timestamp = str(datetime.now())
-            total_fp_elapsed = time() - start_final_pass_time
-            logger.info(
-                f"Finished final-pass (connecting single breaks) at: {finish_fp_timestamp}"
-            )
-            logger.info(f"Total runtime: {total_fp_elapsed} secs")
+        if tracking:
+            lfs = [x for x in output]
+            if tracking_clean_instance_count > 0:
+                lfs = cull_instances(
+                    lfs, tracking_clean_instance_count, tracking_clean_iou_threshold
+                )
+                if not post_connect_single_breaks:
+                    corrected_lfs = connect_single_breaks(
+                        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(
+                    f"Started final-pass (connecting single breaks) at: {start_fp_timestamp}"
+                )
+                corrected_lfs = connect_single_breaks(
+                    lfs, max_instances=tracking_target_instance_count
+                )
+                finish_fp_timestamp = str(datetime.now())
+                total_fp_elapsed = time() - start_final_pass_time
+                logger.info(
+                    f"Finished final-pass (connecting single breaks) at: {finish_fp_timestamp}"
+                )
+                logger.info(f"Total runtime: {total_fp_elapsed} secs")
+            else:
+                corrected_lfs = lfs
 
             output = sio.Labels(
                 labeled_frames=corrected_lfs,
@@ -455,6 +489,9 @@ def run_inference(
             f"Total runtime: {total_elapsed} secs"
         )  # TODO: add number of predicted frames
 
+    if no_empty_frames:
+        output.clean(frames=True, skeletons=False)
+
     if make_labels:
         if output_path is None:
             output_path = Path(

sleap_nn/tracking/tracker.py

@@ -28,6 +28,8 @@
     compute_euclidean_distance,
     compute_iou,
     compute_cosine_sim,
+    cull_instances,
+    cull_frame_instances,
 )
 
 
@@ -74,6 +76,9 @@ class Tracker:
     robust_best_instance: float = 1.0
     use_flow: bool = False
     is_local_queue: bool = False
+    tracking_target_instance_count: int = 0
+    tracking_pre_cull_to_target: int = 0
+    tracking_pre_cull_iou_threshold: float = 0
     _scoring_functions: Dict[str, Any] = {
         "oks": compute_oks,
         "iou": compute_iou,
@@ -114,6 +119,9 @@ 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_pre_cull_to_target: int = 0,
+        tracking_pre_cull_iou_threshold: float = 0,
     ):
         """Create `Tracker` from config.
 
@@ -154,6 +162,9 @@ 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_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)
 
         """
         if candidates_method == "fixed_window":
@@ -189,6 +200,9 @@ def from_config(
                 of_window_size=of_window_size,
                 of_max_levels=of_max_levels,
                 is_local_queue=is_local_queue,
+                tracking_target_instance_count=tracking_target_instance_count,
+                tracking_pre_cull_to_target=tracking_pre_cull_to_target,
+                tracking_pre_cull_iou_threshold=tracking_pre_cull_iou_threshold,
             )
 
         tracker = cls(
@@ -201,6 +215,9 @@ def from_config(
             track_matching_method=track_matching_method,
             use_flow=use_flow,
             is_local_queue=is_local_queue,
+            tracking_target_instance_count=tracking_target_instance_count,
+            tracking_pre_cull_to_target=tracking_pre_cull_to_target,
+            tracking_pre_cull_iou_threshold=tracking_pre_cull_iou_threshold,
         )
         return tracker
 
@@ -220,6 +237,12 @@ 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:
+            untracked_instances = cull_frame_instances(
+                untracked_instances,
+                self.tracking_target_instance_count,
+                self.tracking_pre_cull_iou_threshold,
+            )
         # get features for the untracked instances.
         current_instances = self.get_features(untracked_instances, frame_idx, image)
 
@@ -469,6 +492,9 @@ 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_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)
 
     """
 
@@ -703,6 +729,11 @@ 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_pre_cull_to_target: int = 0,
+    tracking_pre_cull_iou_threshold: float = 0,
+    tracking_clean_instance_count: int = 0,
+    tracking_clean_iou_threshold: float = 0,
 ) -> List[sio.LabeledFrame]:
     """Run tracking on a given set of frames.
 
@@ -746,6 +777,11 @@ 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_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)
+        tracking_clean_iou_threshold: IOU to use when culling instances *after* tracking. (default: 0)
 
     Returns:
         `sio.Labels` object with tracked instances.
@@ -766,6 +802,9 @@ def run_tracker(
         of_img_scale=of_img_scale,
         of_window_size=of_window_size,
         of_max_levels=of_max_levels,
+        tracking_target_instance_count=tracking_target_instance_count,
+        tracking_pre_cull_to_target=tracking_pre_cull_to_target,
+        tracking_pre_cull_iou_threshold=tracking_pre_cull_iou_threshold,
     )
     tracked_lfs = []
     for lf in untracked_frames:
@@ -791,17 +830,28 @@ def run_tracker(
             )
         )
 
+    if tracking_clean_instance_count > 0:
+        tracked_lfs = cull_instances(
+            tracked_lfs, tracking_clean_instance_count, tracking_clean_iou_threshold
+        )
+        if not post_connect_single_breaks:
+            tracked_lfs = connect_single_breaks(
+                tracked_lfs, tracking_clean_instance_count
+            )
+
     if post_connect_single_breaks:
-        if max_tracks is None:
-            message = "Max_tracks is None. To connect single breaks, max_tracks should be set to an integer."
+        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(
             f"Started final-pass (connecting single breaks) at: {start_fp_timestamp}"
         )
-        tracked_lfs = connect_single_breaks(tracked_lfs, max_instances=max_tracks)
+        tracked_lfs = connect_single_breaks(
+            tracked_lfs, max_instances=tracking_target_instance_count
+        )
         finish_fp_timestamp = str(datetime.now())
         total_fp_elapsed = time() - start_final_pass_time
         logger.info(