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

sleap_nn/__init__.py

@@ -41,16 +41,14 @@ def _safe_print(msg):
 
 
 # Add logger with the custom filter
-# Disable colorization to avoid ANSI codes in captured output
 logger.add(
     _safe_print,
     level="DEBUG",
     filter=_should_log,
-    format="{time:YYYY-MM-DD HH:mm:ss} | {message}",
-    colorize=False,
+    format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {name}:{function}:{line} | {message}",
 )
 
-__version__ = "0.1.0"
+__version__ = "0.1.0a0"
 
 # Public API
 from sleap_nn.evaluation import load_metrics

sleap_nn/architectures/convnext.py

@@ -281,10 +281,6 @@ 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,
@@ -297,7 +293,6 @@ 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, Optional, Text, Tuple, Union
+from typing import List, Text, Tuple, Union
 from collections import OrderedDict
 import torch
 from torch import nn
@@ -391,18 +391,10 @@ 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
@@ -477,13 +469,13 @@ class SimpleUpsamplingBlock(nn.Module):
                     first_conv_in_channels = refine_convs_filters
                 else:
                     if self.up_interpolate:
-                        # 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
+                        # 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
                     else:
-                        # With transpose conv, input is transpose_conv_output + skip_channels
+                        # With transpose conv, input is transpose_conv_output + feature channels
                         first_conv_in_channels = (
-                            self.skip_channels + transpose_convs_filters
+                            refine_convs_filters + transpose_convs_filters
                         )
             else:
                 if not self.feat_concat:
@@ -590,7 +582,6 @@ 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__()
@@ -607,7 +598,6 @@ 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
@@ -634,13 +624,6 @@ 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
@@ -659,7 +642,6 @@ 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:
@@ -677,7 +659,6 @@ 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,13 +309,6 @@ 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,
@@ -328,7 +321,6 @@ 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,55 +1,17 @@
-"""Unified CLI for SLEAP-NN using rich-click for styled output."""
+"""Unified CLI for SLEAP-NN using Click."""
 
-import subprocess
-import tempfile
-import shutil
-from datetime import datetime
-
-import rich_click as click
-from click import Command
+import click
 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__
-from sleap_nn.config.utils import get_model_type_from_cfg
 import hydra
 import sys
-
-# 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
+from click import Command
 
 
 def print_version(ctx, param, value):
@@ -93,165 +55,47 @@ def cli():
 
     Use subcommands to run different workflows:
 
-    train    - Run training workflow (auto-handles multi-GPU)
-    track    - Run inference/tracking workflow
+    train    - Run training workflow
+    track    - Run inference/ tracking workflow
     eval     - Run evaluation workflow
     system   - Display system information and GPU status
     """
     pass
 
 
-def _get_num_devices_from_config(cfg: DictConfig) -> int:
-    """Determine the number of devices from config.
-
-    User preferences take precedence over auto-detection:
-    - trainer_device_indices=[0] → 1 device (user choice)
-    - trainer_devices=1 → 1 device (user choice)
-    - trainer_devices="auto" or unset → auto-detect available GPUs
-
-    Returns:
-        Number of devices to use for training.
-    """
-    import torch
-
-    # User preference: explicit device indices (highest priority)
-    device_indices = OmegaConf.select(
-        cfg, "trainer_config.trainer_device_indices", default=None
-    )
-    if device_indices is not None and len(device_indices) > 0:
-        return len(device_indices)
-
-    # User preference: explicit device count
-    devices = OmegaConf.select(cfg, "trainer_config.trainer_devices", default="auto")
-
-    if isinstance(devices, int):
-        return devices
-
-    # Auto-detect only when user hasn't specified (devices is "auto" or None)
-    if devices in ("auto", None, "None"):
-        accelerator = OmegaConf.select(
-            cfg, "trainer_config.trainer_accelerator", default="auto"
-        )
-
-        if accelerator == "cpu":
-            return 1
-        elif torch.cuda.is_available():
-            return torch.cuda.device_count()
-        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-            return 1
-        else:
-            return 1
-
-    return 1
-
-
-def _finalize_config(cfg: DictConfig) -> DictConfig:
-    """Finalize configuration by generating run_name if not provided.
-
-    This runs ONCE before subprocess, ensuring all workers get the same run_name.
-    """
-    # Resolve ckpt_dir first
-    ckpt_dir = OmegaConf.select(cfg, "trainer_config.ckpt_dir", default=None)
-    if ckpt_dir is None or ckpt_dir == "" or ckpt_dir == "None":
-        cfg.trainer_config.ckpt_dir = "."
-
-    # Generate run_name if not provided
-    run_name = OmegaConf.select(cfg, "trainer_config.run_name", default=None)
-    if run_name is None or run_name == "" or run_name == "None":
-        # Get model type from config
-        model_type = get_model_type_from_cfg(cfg)
-
-        # Count frames from labels
-        train_paths = cfg.data_config.train_labels_path
-        val_paths = OmegaConf.select(cfg, "data_config.val_labels_path", default=None)
-
-        train_count = sum(len(sio.load_slp(p)) for p in train_paths)
-        val_count = 0
-        if val_paths:
-            val_count = sum(len(sio.load_slp(p)) for p in val_paths)
-
-        # Generate full run_name with timestamp
-        timestamp = datetime.now().strftime("%y%m%d_%H%M%S")
-        run_name = f"{timestamp}.{model_type}.n={train_count + val_count}"
-        cfg.trainer_config.run_name = run_name
-
-        logger.info(f"Generated run_name: {run_name}")
-
-    return cfg
-
-
 def show_training_help():
-    """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/
+    """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/.
 """
-    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",
-        )
-    )
+    click.echo(help_text)
 
 
 @cli.command(cls=TrainCommand)
+@click.option("--config-name", "-c", type=str, help="Configuration file name")
 @click.option(
-    "--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)"
+    "--config-dir", "-d", type=str, default=".", help="Configuration directory path"
 )
 @click.option(
     "--video-paths",
@@ -283,80 +127,26 @@ sleap-nn train --config-dir /path/to/dir --config-name myrun
     "Can be specified multiple times. "
     'Example: --prefix-map "/old/server/path" "/new/local/path"',
 )
-@click.option(
-    "--video-config",
-    type=str,
-    hidden=True,
-    help="Path to video replacement config YAML (internal use for multi-GPU).",
-)
 @click.argument("overrides", nargs=-1, type=click.UNPROCESSED)
-def train(
-    config,
-    config_name,
-    config_dir,
-    video_paths,
-    video_path_map,
-    prefix_map,
-    video_config,
-    overrides,
-):
+def train(config_name, config_dir, video_paths, video_path_map, prefix_map, overrides):
     """Run training workflow with Hydra config overrides.
 
-    Automatically detects multi-GPU setups and handles run_name synchronization
-    by spawning training in a subprocess with a pre-generated config.
-
     Examples:
-        sleap-nn train path/to/config.yaml
-        sleap-nn train --config path/to/config.yaml trainer_config.max_epochs=100
-        sleap-nn train config.yaml trainer_config.trainer_devices=4
+        sleap-nn train --config-name myconfig --config-dir /path/to/config_dir/
+        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
     """
-    # 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 help if no config name provided
+    if not config_name:
         show_training_help()
         return
 
-    # Check video path options early
-    # If --video-config is provided (from subprocess), load from file
-    if video_config:
-        video_cfg = OmegaConf.load(video_config)
-        video_paths = tuple(video_cfg.video_paths) if video_cfg.video_paths else ()
-        video_path_map = (
-            dict(video_cfg.video_path_map) if video_cfg.video_path_map else None
-        )
-        prefix_map = dict(video_cfg.prefix_map) if video_cfg.prefix_map else None
-
-    has_video_paths = len(video_paths) > 0
-    has_video_path_map = video_path_map is not None
-    has_prefix_map = prefix_map is not None
-    options_used = sum([has_video_paths, has_video_path_map, has_prefix_map])
-
-    if options_used > 1:
-        raise click.UsageError(
-            "Cannot use multiple path replacement options. "
-            "Choose one of: --video-paths, --video-path-map, or --prefix-map."
-        )
-
-    # Load config to detect device count
+    # Initialize Hydra manually
+    # resolve the path to the config directory (hydra expects absolute path)
+    config_dir = Path(config_dir).resolve().as_posix()
     with hydra.initialize_config_dir(config_dir=config_dir, version_base=None):
-        cfg = hydra.compose(config_name=config_name, overrides=overrides)
+        # Compose config with overrides
+        cfg = hydra.compose(config_name=config_name, overrides=list(overrides))
 
         # Validate config
         if not hasattr(cfg, "model_config") or not cfg.model_config:
@@ -365,76 +155,6 @@ def train(
             )
             raise click.Abort()
 
-        num_devices = _get_num_devices_from_config(cfg)
-
-        # Check if run_name is already set (means we're a subprocess or user provided it)
-        run_name = OmegaConf.select(cfg, "trainer_config.run_name", default=None)
-        run_name_is_set = run_name is not None and run_name != "" and run_name != "None"
-
-    # Multi-GPU path: spawn subprocess with finalized config
-    # Only do this if run_name is NOT set (otherwise we'd loop infinitely or user set it)
-    if num_devices > 1 and not run_name_is_set:
-        logger.info(
-            f"Detected {num_devices} devices, using subprocess for run_name sync..."
-        )
-
-        # Load and finalize config (generate run_name, apply overrides)
-        with hydra.initialize_config_dir(config_dir=config_dir, version_base=None):
-            cfg = hydra.compose(config_name=config_name, overrides=overrides)
-            cfg = _finalize_config(cfg)
-
-        # Save finalized config to temp file
-        temp_dir = tempfile.mkdtemp(prefix="sleap_nn_train_")
-        temp_config_path = Path(temp_dir) / "training_config.yaml"
-        OmegaConf.save(cfg, temp_config_path)
-        logger.info(f"Saved finalized config to: {temp_config_path}")
-
-        # Save video replacement config if needed (so subprocess doesn't need CLI args)
-        temp_video_config_path = None
-        if options_used == 1:
-            video_replacement_config = {
-                "video_paths": list(video_paths) if has_video_paths else None,
-                "video_path_map": dict(video_path_map) if has_video_path_map else None,
-                "prefix_map": dict(prefix_map) if has_prefix_map else None,
-            }
-            temp_video_config_path = Path(temp_dir) / "video_replacement.yaml"
-            OmegaConf.save(
-                OmegaConf.create(video_replacement_config), temp_video_config_path
-            )
-            logger.info(f"Saved video replacement config to: {temp_video_config_path}")
-
-        # Build subprocess command (no video args - they're in the temp file)
-        cmd = [sys.executable, "-m", "sleap_nn.cli", "train", str(temp_config_path)]
-        if temp_video_config_path:
-            cmd.extend(["--video-config", str(temp_video_config_path)])
-
-        logger.info(f"Launching subprocess: {' '.join(cmd)}")
-
-        try:
-            process = subprocess.Popen(cmd)
-            result = process.wait()
-            if result != 0:
-                logger.error(f"Training failed with exit code {result}")
-                sys.exit(result)
-        except KeyboardInterrupt:
-            logger.info("Training interrupted, terminating subprocess...")
-            process.terminate()
-            try:
-                process.wait(timeout=5)
-            except subprocess.TimeoutExpired:
-                process.kill()
-                process.wait()
-            sys.exit(1)
-        finally:
-            shutil.rmtree(temp_dir, ignore_errors=True)
-            logger.info("Cleaned up temporary files")
-
-        return
-
-    # Single GPU (or subprocess worker): run directly
-    with hydra.initialize_config_dir(config_dir=config_dir, version_base=None):
-        cfg = hydra.compose(config_name=config_name, overrides=overrides)
-
         logger.info("Input config:")
         logger.info("\n" + OmegaConf.to_yaml(cfg))
 
@@ -442,6 +162,19 @@ def train(
         train_labels = None
         val_labels = None
 
+        # Check that only one replacement option is used
+        # video_paths is a tuple (empty if not used), others are None or dict
+        has_video_paths = len(video_paths) > 0
+        has_video_path_map = video_path_map is not None
+        has_prefix_map = prefix_map is not None
+        options_used = sum([has_video_paths, has_video_path_map, has_prefix_map])
+
+        if options_used > 1:
+            raise click.UsageError(
+                "Cannot use multiple path replacement options. "
+                "Choose one of: --video-paths, --video-path-map, or --prefix-map."
+            )
+
         if options_used == 1:
             # Load train labels
             train_labels = [
@@ -459,12 +192,15 @@ def train(
 
             # Build replacement arguments based on option used
             if has_video_paths:
+                # List of paths (order must match videos in labels file)
                 replace_kwargs = {
                     "new_filenames": [Path(p).as_posix() for p in video_paths]
                 }
             elif has_video_path_map:
+                # Dictionary mapping old filenames to new filenames
                 replace_kwargs = {"filename_map": video_path_map}
             else:  # has_prefix_map
+                # Dictionary mapping old prefixes to new prefixes
                 replace_kwargs = {"prefix_map": prefix_map}
 
             # Apply replacement to train labels
@@ -666,7 +402,7 @@ def train(
 @click.option(
     "--queue_maxsize",
     type=int,
-    default=32,
+    default=8,
     help="Maximum size of the frame buffer queue.",
 )
 @click.option(
@@ -681,36 +417,6 @@ def train(
     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,
@@ -843,12 +549,6 @@ def train(
     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
@@ -913,9 +613,5 @@ def system():
     print_system_info()
 
 
-cli.add_command(export_command)
-cli.add_command(predict_command)
-
-
 if __name__ == "__main__":
     cli()