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

sleap_nn/__init__.py

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

sleap_nn/architectures/convnext.py

@@ -281,6 +281,10 @@ class ConvNextWrapper(nn.Module):
             # Keep the block output filters the same
             x_in_shape = int(self.arch["channels"][-1] * filters_rate)
 
+        # Encoder channels for skip connections (reversed to match decoder order)
+        # The forward pass uses enc_output[::2][::-1] for skip features
+        encoder_channels = self.arch["channels"][::-1]
+
         self.dec = Decoder(
             x_in_shape=x_in_shape,
             current_stride=self.current_stride,
@@ -293,6 +297,7 @@ class ConvNextWrapper(nn.Module):
             block_contraction=self.block_contraction,
             output_stride=self.output_stride,
             up_interpolate=up_interpolate,
+            encoder_channels=encoder_channels,
         )
 
         if len(self.dec.decoder_stack):

sleap_nn/architectures/encoder_decoder.py

@@ -25,7 +25,7 @@ classes.
 See the `EncoderDecoder` base class for requirements for creating new architectures.
 """
 
-from typing import List, Text, Tuple, Union
+from typing import List, Optional, Text, Tuple, Union
 from collections import OrderedDict
 import torch
 from torch import nn
@@ -391,10 +391,18 @@ class SimpleUpsamplingBlock(nn.Module):
         transpose_convs_activation: Text = "relu",
         feat_concat: bool = True,
         prefix: Text = "",
+        skip_channels: Optional[int] = None,
     ) -> None:
         """Initialize the class."""
         super().__init__()
 
+        # Determine skip connection channels
+        # If skip_channels is provided, use it; otherwise fall back to refine_convs_filters
+        # This allows ConvNext/SwinT to specify actual encoder channels
+        self.skip_channels = (
+            skip_channels if skip_channels is not None else refine_convs_filters
+        )
+
         self.x_in_shape = x_in_shape
         self.current_stride = current_stride
         self.upsampling_stride = upsampling_stride
@@ -469,13 +477,13 @@ class SimpleUpsamplingBlock(nn.Module):
                     first_conv_in_channels = refine_convs_filters
                 else:
                     if self.up_interpolate:
-                        # With interpolation, input is x_in_shape + feature channels
-                        # The feature channels are the same as x_in_shape since they come from the same level
-                        first_conv_in_channels = x_in_shape + refine_convs_filters
+                        # With interpolation, input is x_in_shape + skip_channels
+                        # skip_channels may differ from refine_convs_filters for ConvNext/SwinT
+                        first_conv_in_channels = x_in_shape + self.skip_channels
                     else:
-                        # With transpose conv, input is transpose_conv_output + feature channels
+                        # With transpose conv, input is transpose_conv_output + skip_channels
                         first_conv_in_channels = (
-                            refine_convs_filters + transpose_convs_filters
+                            self.skip_channels + transpose_convs_filters
                         )
             else:
                 if not self.feat_concat:
@@ -582,6 +590,7 @@ class Decoder(nn.Module):
         block_contraction: bool = False,
         up_interpolate: bool = True,
         prefix: str = "dec",
+        encoder_channels: Optional[List[int]] = None,
     ) -> None:
         """Initialize the class."""
         super().__init__()
@@ -598,6 +607,7 @@ class Decoder(nn.Module):
         self.block_contraction = block_contraction
         self.prefix = prefix
         self.stride_to_filters = {}
+        self.encoder_channels = encoder_channels
 
         self.current_strides = []
         self.residuals = 0
@@ -624,6 +634,13 @@ class Decoder(nn.Module):
 
             next_stride = current_stride // 2
 
+            # Determine skip channels for this decoder block
+            # If encoder_channels provided, use actual encoder channels
+            # Otherwise fall back to computed filters (for UNet compatibility)
+            skip_channels = None
+            if encoder_channels is not None and block < len(encoder_channels):
+                skip_channels = encoder_channels[block]
+
             if self.stem_blocks > 0 and block >= down_blocks + self.stem_blocks:
                 # This accounts for the case where we dont have any more down block features to concatenate with.
                 # In this case, add a simple upsampling block with a conv layer and with no concatenation
@@ -642,6 +659,7 @@ class Decoder(nn.Module):
                         transpose_convs_batch_norm=False,
                         feat_concat=False,
                         prefix=f"{self.prefix}{block}_s{current_stride}_to_s{next_stride}",
+                        skip_channels=skip_channels,
                     )
                 )
             else:
@@ -659,6 +677,7 @@ class Decoder(nn.Module):
                         transpose_convs_filters=block_filters_out,
                         transpose_convs_batch_norm=False,
                         prefix=f"{self.prefix}{block}_s{current_stride}_to_s{next_stride}",
+                        skip_channels=skip_channels,
                     )
                 )
 

sleap_nn/architectures/swint.py

@@ -309,6 +309,13 @@ class SwinTWrapper(nn.Module):
             self.stem_patch_stride * (2**3) * 2
         )  # stem_stride * down_blocks_stride * final_max_pool_stride
 
+        # Encoder channels for skip connections (reversed to match decoder order)
+        # SwinT channels: embed * 2^i for each stage i, then reversed
+        num_stages = len(self.arch["depths"])
+        encoder_channels = [
+            self.arch["embed"] * (2 ** (num_stages - 1 - i)) for i in range(num_stages)
+        ]
+
         self.dec = Decoder(
             x_in_shape=block_filters,
             current_stride=self.current_stride,
@@ -321,6 +328,7 @@ class SwinTWrapper(nn.Module):
             block_contraction=self.block_contraction,
             output_stride=output_stride,
             up_interpolate=up_interpolate,
+            encoder_channels=encoder_channels,
         )
 
         if len(self.dec.decoder_stack):

sleap_nn/cli.py

@@ -1,17 +1,49 @@
-"""Unified CLI for SLEAP-NN using Click."""
+"""Unified CLI for SLEAP-NN using rich-click for styled output."""
 
-import click
+import rich_click as click
+from click import Command
 from loguru import logger
 from pathlib import Path
 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
 import sys
-from click import Command
+
+# Rich-click configuration for styled help
+click.rich_click.TEXT_MARKUP = "markdown"
+click.rich_click.SHOW_ARGUMENTS = True
+click.rich_click.GROUP_ARGUMENTS_OPTIONS = True
+click.rich_click.STYLE_ERRORS_SUGGESTION = "magenta italic"
+click.rich_click.ERRORS_EPILOGUE = (
+    "Try 'sleap-nn [COMMAND] --help' for more information."
+)
+
+
+def is_config_path(arg: str) -> bool:
+    """Check if an argument looks like a config file path.
+
+    Returns True if the arg ends with .yaml or .yml.
+    """
+    return arg.endswith(".yaml") or arg.endswith(".yml")
+
+
+def split_config_path(config_path: str) -> tuple:
+    """Split a full config path into (config_dir, config_name).
+
+    Args:
+        config_path: Full path to a config file.
+
+    Returns:
+        Tuple of (config_dir, config_name) where config_dir is an absolute path.
+    """
+    path = Path(config_path).resolve()
+    return path.parent.as_posix(), path.name
 
 
 def print_version(ctx, param, value):
@@ -64,38 +96,77 @@ def cli():
 
 
 def show_training_help():
-    """Display training help information."""
-    help_text = """
-sleap-nn train — Train SLEAP models from a config YAML file.
-
-Usage:
-  sleap-nn train --config-dir <dir> --config-name <name> [overrides]
-
-Common overrides:
-  trainer_config.max_epochs=100
-  trainer_config.batch_size=32
-
-Examples:
-  Start new run:
-    sleap-nn train --config-dir /path/to/config_dir/ --config-name myrun
-  Resume 20 more epochs:
-    sleap-nn train --config-dir /path/to/config_dir/ --config-name myrun \\
-      trainer_config.resume_ckpt_path=<path/to/ckpt> \\
-      trainer_config.max_epochs=20
-
-Tips:
-  - Use -m/--multirun for sweeps; outputs go under hydra.sweep.dir.
-  - For Hydra flags and completion, use --hydra-help.
-
-For a detailed list of all available config options, please refer to https://nn.sleap.ai/config/.
+    """Display training help information with rich formatting."""
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.markdown import Markdown
+
+    console = Console()
+
+    help_md = """
+## Usage
+
+```
+sleap-nn train <config.yaml> [overrides]
+sleap-nn train --config <path/to/config.yaml> [overrides]
+```
+
+## Common Overrides
+
+| Override | Description |
+|----------|-------------|
+| `trainer_config.max_epochs=100` | Set maximum training epochs |
+| `trainer_config.batch_size=32` | Set batch size |
+| `trainer_config.save_ckpt=true` | Enable checkpoint saving |
+
+## Examples
+
+**Start a new training run:**
+```bash
+sleap-nn train path/to/config.yaml
+sleap-nn train --config path/to/config.yaml
+```
+
+**With overrides:**
+```bash
+sleap-nn train config.yaml trainer_config.max_epochs=100
+```
+
+**Resume training:**
+```bash
+sleap-nn train config.yaml trainer_config.resume_ckpt_path=/path/to/ckpt
+```
+
+**Legacy usage (still supported):**
+```bash
+sleap-nn train --config-dir /path/to/dir --config-name myrun
+```
+
+## Tips
+
+- Use `-m/--multirun` for sweeps; outputs go under `hydra.sweep.dir`
+- For Hydra flags and completion, use `--hydra-help`
+- Config documentation: https://nn.sleap.ai/config/
 """
-    click.echo(help_text)
+    console.print(
+        Panel(
+            Markdown(help_md),
+            title="[bold cyan]sleap-nn train[/bold cyan]",
+            subtitle="Train SLEAP models from a config YAML file",
+            border_style="cyan",
+        )
+    )
 
 
 @cli.command(cls=TrainCommand)
-@click.option("--config-name", "-c", type=str, help="Configuration file name")
 @click.option(
-    "--config-dir", "-d", type=str, default=".", help="Configuration directory path"
+    "--config",
+    type=str,
+    help="Path to configuration file (e.g., path/to/config.yaml)",
+)
+@click.option("--config-name", "-c", type=str, help="Configuration file name (legacy)")
+@click.option(
+    "--config-dir", "-d", type=str, default=".", help="Configuration directory (legacy)"
 )
 @click.option(
     "--video-paths",
@@ -128,25 +199,43 @@ For a detailed list of all available config options, please refer to https://nn.
     'Example: --prefix-map "/old/server/path" "/new/local/path"',
 )
 @click.argument("overrides", nargs=-1, type=click.UNPROCESSED)
-def train(config_name, config_dir, video_paths, video_path_map, prefix_map, overrides):
+def train(
+    config, config_name, config_dir, video_paths, video_path_map, prefix_map, overrides
+):
     """Run training workflow with Hydra config overrides.
 
     Examples:
-        sleap-nn train --config-name myconfig --config-dir /path/to/config_dir/
+        sleap-nn train path/to/config.yaml
+        sleap-nn train --config path/to/config.yaml trainer_config.max_epochs=100
         sleap-nn train -c myconfig -d /path/to/config_dir/ trainer_config.max_epochs=100
-        sleap-nn train -c myconfig -d /path/to/config_dir/ +experiment=new_model
     """
-    # Show help if no config name provided
-    if not config_name:
+    # Convert overrides to a mutable list
+    overrides = list(overrides)
+
+    # Check if the first positional arg is a config path (not a Hydra override)
+    config_from_positional = None
+    if overrides and is_config_path(overrides[0]):
+        config_from_positional = overrides.pop(0)
+
+    # Resolve config path with priority:
+    # 1. Positional config path (e.g., sleap-nn train config.yaml)
+    # 2. --config flag (e.g., sleap-nn train --config config.yaml)
+    # 3. Legacy --config-dir/--config-name flags
+    if config_from_positional:
+        config_dir, config_name = split_config_path(config_from_positional)
+    elif config:
+        config_dir, config_name = split_config_path(config)
+    elif config_name:
+        config_dir = Path(config_dir).resolve().as_posix()
+    else:
+        # No config provided - show help
         show_training_help()
         return
 
-    # Initialize Hydra manually
-    # resolve the path to the config directory (hydra expects absolute path)
-    config_dir = Path(config_dir).resolve().as_posix()
+    # Initialize Hydra manually (config_dir is already an absolute path)
     with hydra.initialize_config_dir(config_dir=config_dir, version_base=None):
         # Compose config with overrides
-        cfg = hydra.compose(config_name=config_name, overrides=list(overrides))
+        cfg = hydra.compose(config_name=config_name, overrides=overrides)
 
         # Validate config
         if not hasattr(cfg, "model_config") or not cfg.model_config:
@@ -417,6 +506,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,
@@ -549,6 +668,12 @@ def train(config_name, config_dir, video_paths, video_path_map, prefix_map, over
     default=0,
     help="IOU to use when culling instances *after* tracking. (default: 0)",
 )
+@click.option(
+    "--gui",
+    is_flag=True,
+    default=False,
+    help="Output JSON progress for GUI integration instead of Rich progress bar.",
+)
 def track(**kwargs):
     """Run Inference and Tracking workflow."""
     # Convert model_paths from tuple to list
@@ -613,5 +738,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",
+]