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

sleap_nn/cli.py

@@ -1,15 +1,63 @@
-"""Unified CLI for SLEAP-NN using Click."""
+"""Unified CLI for SLEAP-NN using rich-click for styled output."""
 
-import click
+import subprocess
+import tempfile
+import shutil
+from datetime import datetime
+
+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__
+from sleap_nn.config.utils import get_model_type_from_cfg
 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):
+    """Print version and exit."""
+    if not value or ctx.resilient_parsing:
+        return
+    click.echo(f"sleap-nn {__version__}")
+    ctx.exit()
 
 
 class TrainCommand(Command):
@@ -20,73 +68,295 @@ class TrainCommand(Command):
         show_training_help()
 
 
+def parse_path_map(ctx, param, value):
+    """Parse (old, new) path pairs into a dictionary for path mapping options."""
+    if not value:
+        return None
+    result = {}
+    for old_path, new_path in value:
+        result[old_path] = Path(new_path).as_posix()
+    return result
+
+
 @click.group()
+@click.option(
+    "--version",
+    "-v",
+    is_flag=True,
+    callback=print_version,
+    expose_value=False,
+    is_eager=True,
+    help="Show version and exit.",
+)
 def cli():
     """SLEAP-NN: Neural network backend for training and inference for animal pose estimation.
 
     Use subcommands to run different workflows:
 
-    train    - Run training workflow
-    track    - Run inference/ tracking workflow
+    train    - Run training workflow (auto-handles multi-GPU)
+    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."""
-    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",
+    "-v",
+    multiple=True,
+    help="Video paths to replace existing paths in the labels file. "
+    "Order must match the order of videos in the labels file. "
+    "Can be specified multiple times. "
+    "Example: --video-paths /path/to/vid1.mp4 --video-paths /path/to/vid2.mp4",
+)
+@click.option(
+    "--video-path-map",
+    nargs=2,
+    multiple=True,
+    callback=parse_path_map,
+    metavar="OLD NEW",
+    help="Map old video path to new path. Takes two arguments: old path and new path. "
+    "Can be specified multiple times. "
+    'Example: --video-path-map "/old/vid.mp`4" "/new/vid.mp4"',
+)
+@click.option(
+    "--prefix-map",
+    nargs=2,
+    multiple=True,
+    callback=parse_path_map,
+    metavar="OLD NEW",
+    help="Map old path prefix to new prefix. Takes two arguments: old prefix and new prefix. "
+    "Updates ALL videos that share the same prefix. Useful when moving data between machines. "
+    "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_name, config_dir, overrides):
+def train(
+    config,
+    config_name,
+    config_dir,
+    video_paths,
+    video_path_map,
+    prefix_map,
+    video_config,
+    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 --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
+        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
     """
-    # 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()
+    # 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
     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:
@@ -95,9 +365,118 @@ def train(config_name, config_dir, overrides):
             )
             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))
-        run_training(cfg)
+
+        # Handle video path replacement options
+        train_labels = None
+        val_labels = None
+
+        if options_used == 1:
+            # Load train labels
+            train_labels = [
+                sio.load_slp(path) for path in cfg.data_config.train_labels_path
+            ]
+
+            # Load val labels if they exist
+            if (
+                cfg.data_config.val_labels_path is not None
+                and len(cfg.data_config.val_labels_path) > 0
+            ):
+                val_labels = [
+                    sio.load_slp(path) for path in cfg.data_config.val_labels_path
+                ]
+
+            # Build replacement arguments based on option used
+            if has_video_paths:
+                replace_kwargs = {
+                    "new_filenames": [Path(p).as_posix() for p in video_paths]
+                }
+            elif has_video_path_map:
+                replace_kwargs = {"filename_map": video_path_map}
+            else:  # has_prefix_map
+                replace_kwargs = {"prefix_map": prefix_map}
+
+            # Apply replacement to train labels
+            for labels in train_labels:
+                labels.replace_filenames(**replace_kwargs)
+
+            # Apply replacement to val labels if they exist
+            if val_labels:
+                for labels in val_labels:
+                    labels.replace_filenames(**replace_kwargs)
+
+        run_training(config=cfg, train_labels=train_labels, val_labels=val_labels)
 
 
 @cli.command()
@@ -209,6 +588,18 @@ def train(config_name, config_dir, overrides):
     default=False,
     help="Only run inference on unlabeled suggested frames when running on labels dataset. This is useful for generating predictions for initialization during labeling.",
 )
+@click.option(
+    "--exclude_user_labeled",
+    is_flag=True,
+    default=False,
+    help="Skip frames that have user-labeled instances. Useful when predicting on entire video but skipping already-labeled frames.",
+)
+@click.option(
+    "--only_predicted_frames",
+    is_flag=True,
+    default=False,
+    help="Only run inference on frames that already have predictions. Requires .slp input file. Useful for re-predicting with a different model.",
+)
 @click.option(
     "--no_empty_frames",
     is_flag=True,
@@ -275,14 +666,14 @@ def train(config_name, config_dir, overrides):
 @click.option(
     "--queue_maxsize",
     type=int,
-    default=8,
+    default=32,
     help="Maximum size of the frame buffer queue.",
 )
 @click.option(
     "--crop_size",
     type=int,
     default=None,
-    help="Crop size. If not provided, the crop size from training_config.yaml is used.",
+    help="Crop size. If not provided, the crop size from training_config.yaml is used. If `input_scale` is provided, then the cropped image will be resized according to `input_scale`.",
 )
 @click.option(
     "--peak_threshold",
@@ -290,6 +681,36 @@ def train(config_name, config_dir, overrides):
     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,
@@ -422,6 +843,12 @@ def train(config_name, config_dir, overrides):
     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
@@ -474,5 +901,21 @@ def eval(**kwargs):
     run_evaluation(**kwargs)
 
 
+@cli.command()
+def system():
+    """Display system information and GPU status.
+
+    Shows Python version, platform, PyTorch version, CUDA availability,
+    driver version with compatibility check, GPU details, and package versions.
+    """
+    from sleap_nn.system_info import print_system_info
+
+    print_system_info()
+
+
+cli.add_command(export_command)
+cli.add_command(predict_command)
+
+
 if __name__ == "__main__":
     cli()