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

sleap_nn/predict.py

@@ -67,16 +67,13 @@ def run_inference(
     only_predicted_frames: bool = False,
     no_empty_frames: bool = False,
     batch_size: int = 4,
-    queue_maxsize: int = 32,
+    queue_maxsize: int = 8,
     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,
@@ -113,7 +110,6 @@ 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.
 
@@ -151,7 +147,7 @@ def run_inference(
         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: 32.
+        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
                 an .slp path as an alternative to specifying the video path.
         video_dataset: (str) The dataset for HDF5 videos.
@@ -164,15 +160,6 @@ def run_inference(
                 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"`.
@@ -263,8 +250,6 @@ 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
@@ -448,6 +433,13 @@ 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
@@ -466,9 +458,6 @@ def run_inference(
             anchor_part=anchor_part,
         )
 
-        # Set GUI mode for progress output
-        predictor.gui = gui
-
         if (
             tracking
             and not isinstance(predictor, BottomUpMultiClassPredictor)
@@ -564,20 +553,6 @@ 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:
@@ -632,9 +607,6 @@ def run_inference(
         # 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,

sleap_nn/train.py

@@ -118,70 +118,6 @@ def run_training(
                 logger.info(f"p90 dist: {metrics['distance_metrics']['p90']}")
                 logger.info(f"p50 dist: {metrics['distance_metrics']['p50']}")
 
-                # Log test metrics to wandb summary
-                if (
-                    d_name.startswith("test")
-                    and trainer.config.trainer_config.use_wandb
-                ):
-                    import wandb
-
-                    if wandb.run is not None:
-                        summary_metrics = {
-                            f"eval/{d_name}/mOKS": metrics["mOKS"]["mOKS"],
-                            f"eval/{d_name}/oks_voc_mAP": metrics["voc_metrics"][
-                                "oks_voc.mAP"
-                            ],
-                            f"eval/{d_name}/oks_voc_mAR": metrics["voc_metrics"][
-                                "oks_voc.mAR"
-                            ],
-                            f"eval/{d_name}/mPCK": metrics["pck_metrics"]["mPCK"],
-                            f"eval/{d_name}/PCK_5": metrics["pck_metrics"]["PCK@5"],
-                            f"eval/{d_name}/PCK_10": metrics["pck_metrics"]["PCK@10"],
-                            f"eval/{d_name}/distance_avg": metrics["distance_metrics"][
-                                "avg"
-                            ],
-                            f"eval/{d_name}/distance_p50": metrics["distance_metrics"][
-                                "p50"
-                            ],
-                            f"eval/{d_name}/distance_p95": metrics["distance_metrics"][
-                                "p95"
-                            ],
-                            f"eval/{d_name}/distance_p99": metrics["distance_metrics"][
-                                "p99"
-                            ],
-                            f"eval/{d_name}/visibility_precision": metrics[
-                                "visibility_metrics"
-                            ]["precision"],
-                            f"eval/{d_name}/visibility_recall": metrics[
-                                "visibility_metrics"
-                            ]["recall"],
-                        }
-                        for key, value in summary_metrics.items():
-                            wandb.run.summary[key] = value
-
-            # Finish wandb run and cleanup after all evaluation is complete
-            if trainer.config.trainer_config.use_wandb:
-                import wandb
-                import shutil
-
-                if wandb.run is not None:
-                    wandb.finish()
-
-                # Delete local wandb logs if configured
-                wandb_config = trainer.config.trainer_config.wandb
-                should_delete_wandb_logs = wandb_config.delete_local_logs is True or (
-                    wandb_config.delete_local_logs is None
-                    and wandb_config.wandb_mode != "offline"
-                )
-                if should_delete_wandb_logs:
-                    wandb_dir = run_path / "wandb"
-                    if wandb_dir.exists():
-                        logger.info(
-                            f"Deleting local wandb logs at {wandb_dir}... "
-                            "(set trainer_config.wandb.delete_local_logs=false to disable)"
-                        )
-                        shutil.rmtree(wandb_dir, ignore_errors=True)
-
 
 def train(
     train_labels_path: Optional[List[str]] = None,
@@ -203,9 +139,9 @@ def train(
     crop_size: Optional[int] = None,
     min_crop_size: Optional[int] = 100,
     crop_padding: Optional[int] = None,
-    use_augmentations_train: bool = True,
+    use_augmentations_train: bool = False,
     intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
-    geometry_aug: Optional[Union[str, List[str], Dict[str, Any]]] = "rotation",
+    geometry_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
     init_weight: str = "default",
     pretrained_backbone_weights: Optional[str] = None,
     pretrained_head_weights: Optional[str] = None,
@@ -239,7 +175,6 @@ def train(
     wandb_save_viz_imgs_wandb: bool = False,
     wandb_resume_prv_runid: Optional[str] = None,
     wandb_group_name: Optional[str] = None,
-    wandb_delete_local_logs: Optional[bool] = None,
     optimizer: str = "Adam",
     learning_rate: float = 1e-3,
     amsgrad: bool = False,
@@ -306,7 +241,7 @@ def train(
             crop size. If `None`, padding is auto-computed based on augmentation settings.
             Only used when `crop_size` is `None`. Default: None.
         use_augmentations_train: True if the data augmentation should be applied to the
-            training data, else False. Default: True.
+            training data, else False. Default: False.
         intensity_aug: One of ["uniform_noise", "gaussian_noise", "contrast", "brightness"]
             or list of strings from the above allowed values. To have custom values, pass
             a dict with the structure in `sleap_nn.config.data_config.IntensityConfig`.
@@ -318,8 +253,7 @@ def train(
             or list of strings from the above allowed values. To have custom values, pass
             a dict with the structure in `sleap_nn.config.data_config.GeometryConfig`.
             For eg: {
-                        "rotation_min": -45,
-                        "rotation_max": 45,
+                        "rotation": 45,
                         "affine_p": 1.0
                     }
         init_weight: model weights initialization method. "default" uses kaiming uniform
@@ -419,9 +353,6 @@ def train(
         wandb_resume_prv_runid: Previous run ID if training should be resumed from a previous
             ckpt. Default: None
         wandb_group_name: Group name for the wandb run. Default: None.
-        wandb_delete_local_logs: If True, delete local wandb logs folder after training.
-            If False, keep the folder. If None (default), automatically delete if logging
-            online (wandb_mode != "offline") and keep if logging offline. Default: None.
         optimizer: Optimizer to be used. One of ["Adam", "AdamW"]. Default: "Adam".
         learning_rate: Learning rate of type float. Default: 1e-3.
         amsgrad: Enable AMSGrad with the optimizer. Default: False.
@@ -525,7 +456,6 @@ def train(
         wandb_save_viz_imgs_wandb=wandb_save_viz_imgs_wandb,
         wandb_resume_prv_runid=wandb_resume_prv_runid,
         wandb_group_name=wandb_group_name,
-        wandb_delete_local_logs=wandb_delete_local_logs,
         optimizer=optimizer,
         learning_rate=learning_rate,
         amsgrad=amsgrad,