sleap-nn 0.0.5__py3-none-any.whl → 0.1.0__py3-none-any.whl

sleap_nn/predict.py

@@ -15,6 +15,11 @@ from sleap_nn.tracking.tracker import (
     connect_single_breaks,
     cull_instances,
 )
+from sleap_nn.system_info import get_startup_info_string
+from sleap_nn.inference.provenance import (
+    build_inference_provenance,
+    build_tracking_only_provenance,
+)
 from omegaconf import OmegaConf
 import sleap_io as sio
 from pathlib import Path
@@ -58,15 +63,20 @@ def run_inference(
     anchor_part: Optional[str] = None,
     only_labeled_frames: bool = False,
     only_suggested_frames: bool = False,
+    exclude_user_labeled: bool = False,
+    only_predicted_frames: bool = False,
     no_empty_frames: bool = False,
     batch_size: int = 4,
-    queue_maxsize: int = 8,
+    queue_maxsize: int = 32,
     video_index: Optional[int] = None,
     video_dataset: Optional[str] = None,
     video_input_format: str = "channels_last",
     frames: Optional[list] = None,
     crop_size: Optional[int] = None,
     peak_threshold: Union[float, List[float]] = 0.2,
+    filter_overlapping: bool = False,
+    filter_overlapping_method: str = "iou",
+    filter_overlapping_threshold: float = 0.8,
     integral_refinement: Optional[str] = "integral",
     integral_patch_size: int = 5,
     return_confmaps: bool = False,
@@ -103,6 +113,7 @@ def run_inference(
     tracking_pre_cull_iou_threshold: float = 0,
     tracking_clean_instance_count: int = 0,
     tracking_clean_iou_threshold: float = 0,
+    gui: bool = False,
 ):
     """Entry point to run inference on trained SLEAP-NN models.
 
@@ -136,21 +147,32 @@ 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`.
+        exclude_user_labeled: (bool) `True` to skip frames that have user-labeled instances. Default: `False`.
+        only_predicted_frames: (bool) `True` to run inference only on frames that already have predictions. 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.
+        queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 32.
         video_index: (int) Integer index of video in .slp file to predict on. To be used with
                 an .slp path as an alternative to specifying the video path.
         video_dataset: (str) The dataset for HDF5 videos.
         video_input_format: (str) The input_format for HDF5 videos.
         frames: (list) List of frames indices. If `None`, all frames in the video are used. Default: None.
         crop_size: (int) Crop size. If not provided, the crop size from training_config.yaml is used.
-                Default: None.
+                If `input_scale` is provided, then the cropped image will be resized according to `input_scale`. Default: None.
         peak_threshold: (float) Minimum confidence threshold. Peaks with values below
                 this will be ignored. Default: 0.2. This can also be `List[float]` for topdown
                 centroid and centered-instance model, where the first element corresponds
                 to centroid model peak finding threshold and the second element is for
                 centered-instance model peak finding.
+        filter_overlapping: (bool) If True, removes overlapping instances after
+                inference using greedy NMS. Applied independently of tracking.
+                Default: False.
+        filter_overlapping_method: (str) Similarity metric for filtering overlapping
+                instances. One of "iou" (bounding box) or "oks" (keypoint similarity).
+                Default: "iou".
+        filter_overlapping_threshold: (float) Similarity threshold for filtering.
+                Instances with similarity > threshold are removed (keeping higher-scoring).
+                Typical values: 0.3 (aggressive) to 0.8 (permissive). Default: 0.8.
         integral_refinement: (str) If `None`, returns the grid-aligned peaks with no refinement.
                 If `"integral"`, peaks will be refined with integral regression.
                 Default: `"integral"`.
@@ -241,6 +263,8 @@ def run_inference(
         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)
+        gui: (bool) If True, outputs JSON progress lines for GUI integration instead
+                of Rich progress bars. Default: False.
 
     Returns:
         Returns `sio.Labels` object if `make_labels` is True. Else this function returns
@@ -256,6 +280,27 @@ def run_inference(
         "scale": input_scale,
     }
 
+    # Validate mutually exclusive frame filter flags
+    if only_labeled_frames and exclude_user_labeled:
+        message = (
+            "--only_labeled_frames and --exclude_user_labeled are mutually exclusive "
+            "(would result in zero frames)"
+        )
+        logger.error(message)
+        raise ValueError(message)
+
+    if (
+        only_predicted_frames
+        and data_path is not None
+        and not data_path.endswith(".slp")
+    ):
+        message = (
+            "--only_predicted_frames requires a .slp file input "
+            "(need Labels to know which frames have predictions)"
+        )
+        logger.error(message)
+        raise ValueError(message)
+
     if model_paths is None or not len(
         model_paths
     ):  # if model paths is not provided, run tracking-only pipeline.
@@ -273,7 +318,8 @@ def run_inference(
                 raise ValueError(message)
 
             start_inf_time = time()
-            start_timestamp = str(datetime.now())
+            start_datetime = datetime.now()
+            start_timestamp = str(start_datetime)
             logger.info(f"Started tracking at: {start_timestamp}")
 
             labels = sio.load_slp(data_path) if input_labels is None else input_labels
@@ -302,7 +348,22 @@ def run_inference(
 
             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."
+                    features_requested = []
+                    if post_connect_single_breaks:
+                        features_requested.append("--post_connect_single_breaks")
+                    if tracking_pre_cull_to_target:
+                        features_requested.append("--tracking_pre_cull_to_target")
+                    features_str = " and ".join(features_requested)
+
+                    if max_tracks is not None:
+                        suggestion = f"Add --tracking_target_instance_count {max_tracks} to your command (using your --max_tracks value)."
+                    else:
+                        suggestion = "Add --tracking_target_instance_count N where N is the expected number of instances per frame."
+
+                    message = (
+                        f"{features_str} requires --tracking_target_instance_count to be set. "
+                        f"{suggestion}"
+                    )
                     logger.error(message)
                     raise ValueError(message)
                 elif tracking_target_instance_count is None:
@@ -332,21 +393,53 @@ def run_inference(
                 tracking_clean_iou_threshold=tracking_clean_iou_threshold,
             )
 
-            finish_timestamp = str(datetime.now())
+            end_datetime = datetime.now()
+            finish_timestamp = str(end_datetime)
             total_elapsed = time() - start_inf_time
             logger.info(f"Finished tracking at: {finish_timestamp}")
             logger.info(f"Total runtime: {total_elapsed} secs")
 
+            # Build tracking-only provenance
+            tracking_params = {
+                "window_size": tracking_window_size,
+                "min_new_track_points": min_new_track_points,
+                "candidates_method": candidates_method,
+                "min_match_points": min_match_points,
+                "features": features,
+                "scoring_method": scoring_method,
+                "scoring_reduction": scoring_reduction,
+                "robust_best_instance": robust_best_instance,
+                "track_matching_method": track_matching_method,
+                "max_tracks": max_tracks,
+                "use_flow": use_flow,
+                "post_connect_single_breaks": post_connect_single_breaks,
+            }
+            provenance = build_tracking_only_provenance(
+                input_labels=labels,
+                input_path=data_path,
+                start_time=start_datetime,
+                end_time=end_datetime,
+                tracking_params=tracking_params,
+                frames_processed=len(tracked_frames),
+            )
+
             output = sio.Labels(
                 labeled_frames=tracked_frames,
                 videos=labels.videos,
                 skeletons=labels.skeletons,
+                provenance=provenance,
             )
 
     else:
         start_inf_time = time()
-        start_timestamp = str(datetime.now())
+        start_datetime = datetime.now()
+        start_timestamp = str(start_datetime)
         logger.info(f"Started inference at: {start_timestamp}")
+        logger.info(get_startup_info_string())
+
+        # Convert device to string if it's a torch.device object
+        if hasattr(device, "type"):
+            device = str(device)
 
         if device == "auto":
             device = (
@@ -355,13 +448,6 @@ def run_inference(
                 else "mps" if torch.backends.mps.is_available() else "cpu"
             )
 
-        if integral_refinement is not None and device == "mps":  # TODO
-            # kornia/geometry/transform/imgwarp.py:382: in get_perspective_transform. NotImplementedError: The operator 'aten::_linalg_solve_ex.result' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
-            logger.info(
-                "Integral refinement is not supported with MPS accelerator. Setting integral refinement to None."
-            )
-            integral_refinement = None
-
         logger.info(f"Using device: {device}")
 
         # initializes the inference model
@@ -380,6 +466,9 @@ def run_inference(
             anchor_part=anchor_part,
         )
 
+        # Set GUI mode for progress output
+        predictor.gui = gui
+
         if (
             tracking
             and not isinstance(predictor, BottomUpMultiClassPredictor)
@@ -387,7 +476,22 @@ def run_inference(
         ):
             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."
+                    features_requested = []
+                    if post_connect_single_breaks:
+                        features_requested.append("--post_connect_single_breaks")
+                    if tracking_pre_cull_to_target:
+                        features_requested.append("--tracking_pre_cull_to_target")
+                    features_str = " and ".join(features_requested)
+
+                    if max_tracks is not None:
+                        suggestion = f"Add --tracking_target_instance_count {max_tracks} to your command (using your --max_tracks value)."
+                    else:
+                        suggestion = "Add --tracking_target_instance_count N or --max_instances N where N is the expected number of instances per frame."
+
+                    message = (
+                        f"{features_str} requires --tracking_target_instance_count or --max_instances to be set. "
+                        f"{suggestion}"
+                    )
                     logger.error(message)
                     raise ValueError(message)
                 elif tracking_target_instance_count is None:
@@ -448,6 +552,8 @@ def run_inference(
             frames=frames,
             only_labeled_frames=only_labeled_frames,
             only_suggested_frames=only_suggested_frames,
+            exclude_user_labeled=exclude_user_labeled,
+            only_predicted_frames=only_predicted_frames,
             video_index=video_index,
             video_dataset=video_dataset,
             video_input_format=video_input_format,
@@ -458,6 +564,20 @@ def run_inference(
             make_labels=make_labels,
         )
 
+        # Filter overlapping instances (independent of tracking)
+        if filter_overlapping and make_labels:
+            from sleap_nn.inference.postprocessing import filter_overlapping_instances
+
+            output = filter_overlapping_instances(
+                output,
+                threshold=filter_overlapping_threshold,
+                method=filter_overlapping_method,
+            )
+            logger.info(
+                f"Filtered overlapping instances with {filter_overlapping_method.upper()} "
+                f"threshold: {filter_overlapping_threshold}"
+            )
+
         if tracking:
             lfs = [x for x in output]
             if tracking_clean_instance_count > 0:
@@ -492,12 +612,97 @@ def run_inference(
                 skeletons=output.skeletons,
             )
 
-        finish_timestamp = str(datetime.now())
+        end_datetime = datetime.now()
+        finish_timestamp = str(end_datetime)
         total_elapsed = time() - start_inf_time
         logger.info(f"Finished inference at: {finish_timestamp}")
-        logger.info(
-            f"Total runtime: {total_elapsed} secs"
-        )  # TODO: add number of predicted frames
+        logger.info(f"Total runtime: {total_elapsed} secs")
+
+        # Determine input labels for provenance preservation
+        input_labels_for_prov = None
+        if input_labels is not None:
+            input_labels_for_prov = input_labels
+        elif data_path is not None and data_path.endswith(".slp"):
+            # Load input labels to preserve provenance (if not already loaded)
+            try:
+                input_labels_for_prov = sio.load_slp(data_path)
+            except Exception:
+                pass
+
+        # Build inference parameters for provenance
+        inference_params = {
+            "peak_threshold": peak_threshold,
+            "filter_overlapping": filter_overlapping,
+            "filter_overlapping_method": filter_overlapping_method,
+            "filter_overlapping_threshold": filter_overlapping_threshold,
+            "integral_refinement": integral_refinement,
+            "integral_patch_size": integral_patch_size,
+            "batch_size": batch_size,
+            "max_instances": max_instances,
+            "crop_size": crop_size,
+            "input_scale": input_scale,
+            "anchor_part": anchor_part,
+        }
+
+        # Build tracking parameters if tracking was enabled
+        tracking_params_prov = None
+        if tracking:
+            tracking_params_prov = {
+                "window_size": tracking_window_size,
+                "min_new_track_points": min_new_track_points,
+                "candidates_method": candidates_method,
+                "min_match_points": min_match_points,
+                "features": features,
+                "scoring_method": scoring_method,
+                "scoring_reduction": scoring_reduction,
+                "robust_best_instance": robust_best_instance,
+                "track_matching_method": track_matching_method,
+                "max_tracks": max_tracks,
+                "use_flow": use_flow,
+                "post_connect_single_breaks": post_connect_single_breaks,
+            }
+
+        # Determine frame selection method
+        frame_selection_method = "all"
+        if only_labeled_frames:
+            frame_selection_method = "labeled"
+        elif only_suggested_frames:
+            frame_selection_method = "suggested"
+        elif only_predicted_frames:
+            frame_selection_method = "predicted"
+        elif frames is not None:
+            frame_selection_method = "specified"
+
+        # Determine model type from predictor class
+        predictor_type_map = {
+            "TopDownPredictor": "top_down",
+            "SingleInstancePredictor": "single_instance",
+            "BottomUpPredictor": "bottom_up",
+            "BottomUpMultiClassPredictor": "bottom_up_multi_class",
+            "TopDownMultiClassPredictor": "top_down_multi_class",
+        }
+        model_type = predictor_type_map.get(type(predictor).__name__)
+
+        # Build and set provenance (only for Labels objects)
+        if make_labels and isinstance(output, sio.Labels):
+            provenance = build_inference_provenance(
+                model_paths=model_paths,
+                model_type=model_type,
+                start_time=start_datetime,
+                end_time=end_datetime,
+                input_labels=input_labels_for_prov,
+                input_path=data_path,
+                frames_processed=(
+                    len(output.labeled_frames)
+                    if hasattr(output, "labeled_frames")
+                    else None
+                ),
+                frame_selection_method=frame_selection_method,
+                inference_params=inference_params,
+                tracking_params=tracking_params_prov,
+                device=device,
+            )
+            output.provenance = provenance
 
     if no_empty_frames:
         output.clean(frames=True, skeletons=False)