sleap-nn 0.1.0a3__tar.gz → 0.1.0a4__tar.gz

{sleap_nn-0.1.0a3 → sleap_nn-0.1.0a4}/.github/workflows/docs.yml

@@ -57,7 +57,7 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: |
-            uv run mike deploy --allow-empty --push "${{ github.event.release.tag_name }}"
+            uv run mike deploy --update-aliases --allow-empty --push "${{ github.event.release.tag_name }}" prerelease
 
       - name: Build and upload docs (dev)
         if: ${{ github.event_name == 'push' }}

{sleap_nn-0.1.0a3 → sleap_nn-0.1.0a4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.1.0a3
+Version: 0.1.0a4
 Summary: Neural network backend for training and inference for animal pose estimation.
 Author-email: Divya Seshadri Murali <dimurali@salk.edu>, Elizabeth Berrigan <eberrigan@salk.edu>, Vincent Tu <vitu@ucsd.edu>, Liezl Maree <lmaree@salk.edu>, David Samy <davidasamy@gmail.com>, Talmo Pereira <talmo@salk.edu>
 License: BSD-3-Clause
@@ -32,6 +32,7 @@ Requires-Dist: hydra-core
 Requires-Dist: jupyter
 Requires-Dist: jupyterlab
 Requires-Dist: pyzmq
+Requires-Dist: rich-click>=1.9.5
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
 Requires-Dist: torchvision>=0.20.0; extra == "torch"

{sleap_nn-0.1.0a3 → sleap_nn-0.1.0a4}/pyproject.toml

@@ -48,6 +48,7 @@ dependencies = [
     "jupyter",
     "jupyterlab",
     "pyzmq",
+    "rich-click>=1.9.5",
 ]
 dynamic = ["version", "readme"]
 

{sleap_nn-0.1.0a3 → sleap_nn-0.1.0a4}/sleap_nn/__init__.py

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

{sleap_nn-0.1.0a3 → sleap_nn-0.1.0a4}/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-0.1.0a3 → sleap_nn-0.1.0a4}/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-0.1.0a3 → sleap_nn-0.1.0a4}/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-0.1.0a3 → sleap_nn-0.1.0a4}/sleap_nn/cli.py

@@ -1,6 +1,7 @@
-"""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
@@ -13,7 +14,36 @@ 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):
@@ -66,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",
@@ -130,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:
@@ -581,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

{sleap_nn-0.1.0a3 → sleap_nn-0.1.0a4}/sleap_nn/inference/peak_finding.py

@@ -3,9 +3,8 @@
 from typing import Optional, Tuple
 
 import kornia as K
-import numpy as np
 import torch
-from kornia.geometry.transform import crop_and_resize
+import torch.nn.functional as F
 
 from sleap_nn.data.instance_cropping import make_centered_bboxes
 
@@ -13,7 +12,11 @@ from sleap_nn.data.instance_cropping import make_centered_bboxes
 def crop_bboxes(
     images: torch.Tensor, bboxes: torch.Tensor, sample_inds: torch.Tensor
 ) -> torch.Tensor:
-    """Crop bounding boxes from a batch of images.
+    """Crop bounding boxes from a batch of images using fast tensor indexing.
+
+    This uses tensor unfold operations to extract patches, which is significantly
+    faster than kornia's crop_and_resize (17-51x speedup) as it avoids perspective
+    transform computations.
 
     Args:
         images: Tensor of shape (samples, channels, height, width) of a batch of images.
@@ -27,7 +30,7 @@ def crop_bboxes(
             box should be cropped from.
 
     Returns:
-        A tensor of shape (n_bboxes, crop_height, crop_width, channels) of the same
+        A tensor of shape (n_bboxes, channels, crop_height, crop_width) of the same
         dtype as the input image. The crop size is inferred from the bounding box
         coordinates.
 
@@ -42,26 +45,53 @@ def crop_bboxes(
 
     See also: `make_centered_bboxes`
     """
+    n_crops = bboxes.shape[0]
+    if n_crops == 0:
+        # Return empty tensor; use default crop size since we can't infer from bboxes
+        return torch.empty(
+            0, images.shape[1], 0, 0, device=images.device, dtype=images.dtype
+        )
+
     # Compute bounding box size to use for crops.
-    height = abs(bboxes[0, 3, 1] - bboxes[0, 0, 1])
-    width = abs(bboxes[0, 1, 0] - bboxes[0, 0, 0])
-    box_size = tuple(torch.round(torch.Tensor((height + 1, width + 1))).to(torch.int32))
+    height = int(abs(bboxes[0, 3, 1] - bboxes[0, 0, 1]).item()) + 1
+    width = int(abs(bboxes[0, 1, 0] - bboxes[0, 0, 0]).item()) + 1
 
     # Store original dtype for conversion back after cropping.
     original_dtype = images.dtype
+    device = images.device
+    n_samples, channels, img_h, img_w = images.shape
+    half_h, half_w = height // 2, width // 2
 
-    # Kornia's crop_and_resize requires float32 input.
-    images_to_crop = images[sample_inds]
-    if not torch.is_floating_point(images_to_crop):
-        images_to_crop = images_to_crop.float()
-
-    # Crop.
-    crops = crop_and_resize(
-        images_to_crop,  # (n_boxes, channels, height, width)
-        boxes=bboxes,
-        size=box_size,
+    # Pad images for edge handling.
+    images_padded = F.pad(
+        images.float(), (half_w, half_w, half_h, half_h), mode="constant", value=0
     )
 
+    # Extract all possible patches using unfold (creates a view, no copy).
+    # Shape after unfold: (n_samples, channels, img_h, img_w, height, width)
+    patches = images_padded.unfold(2, height, 1).unfold(3, width, 1)
+
+    # Get crop centers from bboxes.
+    # The bbox top-left is at index 0, with (x, y) coordinates.
+    # We need the center of the crop (peak location), which is top-left + half_size.
+    # Ensure bboxes are on the same device as images for index computation.
+    bboxes_on_device = bboxes.to(device)
+    crop_x = (bboxes_on_device[:, 0, 0] + half_w).to(torch.long)
+    crop_y = (bboxes_on_device[:, 0, 1] + half_h).to(torch.long)
+
+    # Clamp indices to valid bounds to handle edge cases where centroids
+    # might be at or beyond image boundaries.
+    crop_x = torch.clamp(crop_x, 0, patches.shape[3] - 1)
+    crop_y = torch.clamp(crop_y, 0, patches.shape[2] - 1)
+
+    # Select crops using advanced indexing.
+    # Convert sample_inds to tensor if it's a list.
+    if not isinstance(sample_inds, torch.Tensor):
+        sample_inds = torch.tensor(sample_inds, device=device)
+    sample_inds_long = sample_inds.to(device=device, dtype=torch.long)
+    crops = patches[sample_inds_long, :, crop_y, crop_x]
+    # Shape: (n_crops, channels, height, width)
+
     # Cast back to original dtype and return.
     crops = crops.to(original_dtype)
     return crops