sleap-nn 0.1.0a2__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

sleap_nn/__init__.py

@@ -50,7 +50,7 @@ logger.add(
     colorize=False,
 )
 
-__version__ = "0.1.0a2"
+__version__ = "0.1.0a3"
 
 # Public API
 from sleap_nn.evaluation import load_metrics

sleap_nn/cli.py

@@ -7,6 +7,8 @@ from omegaconf import OmegaConf, DictConfig
 import sleap_io as sio
 from sleap_nn.predict import run_inference, frame_list
 from sleap_nn.evaluation import run_evaluation
+from sleap_nn.export.cli import export as export_command
+from sleap_nn.export.cli import predict as predict_command
 from sleap_nn.train import run_training
 from sleap_nn import __version__
 import hydra
@@ -417,6 +419,36 @@ def train(config_name, config_dir, video_paths, video_path_map, prefix_map, over
     default=0.2,
     help="Minimum confidence map value to consider a peak as valid.",
 )
+@click.option(
+    "--filter_overlapping",
+    is_flag=True,
+    default=False,
+    help=(
+        "Enable filtering of overlapping instances after inference using greedy NMS. "
+        "Applied independently of tracking. (default: False)"
+    ),
+)
+@click.option(
+    "--filter_overlapping_method",
+    type=click.Choice(["iou", "oks"]),
+    default="iou",
+    help=(
+        "Similarity metric for filtering overlapping instances. "
+        "'iou': bounding box intersection-over-union. "
+        "'oks': Object Keypoint Similarity (pose-based). (default: iou)"
+    ),
+)
+@click.option(
+    "--filter_overlapping_threshold",
+    type=float,
+    default=0.8,
+    help=(
+        "Similarity threshold for filtering overlapping instances. "
+        "Instances with similarity above this threshold are removed, "
+        "keeping the higher-scoring instance. "
+        "Typical values: 0.3 (aggressive) to 0.8 (permissive). (default: 0.8)"
+    ),
+)
 @click.option(
     "--integral_refinement",
     type=str,
@@ -613,5 +645,9 @@ def system():
     print_system_info()
 
 
+cli.add_command(export_command)
+cli.add_command(predict_command)
+
+
 if __name__ == "__main__":
     cli()

sleap_nn/evaluation.py

@@ -639,11 +639,19 @@ class Evaluator:
         mPCK_parts = pcks.mean(axis=0).mean(axis=-1)
         mPCK = mPCK_parts.mean()
 
+        # Precompute PCK at common thresholds
+        idx_5 = np.argmin(np.abs(thresholds - 5))
+        idx_10 = np.argmin(np.abs(thresholds - 10))
+        pck5 = pcks[:, :, idx_5].mean()
+        pck10 = pcks[:, :, idx_10].mean()
+
         return {
             "thresholds": thresholds,
             "pcks": pcks,
             "mPCK_parts": mPCK_parts,
             "mPCK": mPCK,
+            "PCK@5": pck5,
+            "PCK@10": pck10,
         }
 
     def visibility_metrics(self):

sleap_nn/export/__init__.py

@@ -0,0 +1,21 @@
+"""Export utilities for sleap-nn."""
+
+from sleap_nn.export.exporters import export_model, export_to_onnx, export_to_tensorrt
+from sleap_nn.export.metadata import ExportMetadata
+from sleap_nn.export.predictors import (
+    load_exported_model,
+    ONNXPredictor,
+    TensorRTPredictor,
+)
+from sleap_nn.export.utils import build_bottomup_candidate_template
+
+__all__ = [
+    "export_model",
+    "export_to_onnx",
+    "export_to_tensorrt",
+    "load_exported_model",
+    "ONNXPredictor",
+    "TensorRTPredictor",
+    "ExportMetadata",
+    "build_bottomup_candidate_template",
+]