cli

sleap_nn.cli

Unified CLI for SLEAP-NN using rich-click for styled output.

Classes:

Name	Description
LazyGroup	Click group that lazily loads export subcommands on first use.
TrainCommand	Custom command class that overrides help behavior for train command.

Functions:

Name	Description
cli	SLEAP-NN: Neural network backend for training and inference for animal pose estimation.
config	Generate training configuration for a SLEAP file.
eval	Run evaluation workflow.
eval_tracking	Evaluate identity persistence of a tracked prediction.
infer	Deprecated alias for predict. Use sleap-nn predict instead.
info	Display model configuration and evaluation metrics.
is_config_path	Check if an argument looks like a config file path.
parse_path_map	Parse (old, new) path pairs into a dictionary for path mapping options.
predict	Run inference on videos or labels files.
print_version	Print version and exit.
show_training_help	Display training help information with rich formatting.
split_config_path	Split a full config path into (config_dir, config_name).
system	Display system information and GPU status.
track	Run Inference and Tracking workflow (legacy pipeline).
train	Run training workflow with Hydra config overrides.

LazyGroup

Bases: RichGroup

Click group that lazily loads export subcommands on first use.

Methods:

Name	Description
get_command	Get a command by name, loading export subcommands on first access.
list_commands	List all commands, loading export subcommands on first access.

Source code in sleap_nn/cli.py

class LazyGroup(click.RichGroup):
    """Click group that lazily loads export subcommands on first use."""

    _export_loaded = False

    def list_commands(self, ctx):
        """List all commands, loading export subcommands on first access."""
        self._ensure_export_loaded()
        return super().list_commands(ctx)

    def get_command(self, ctx, cmd_name):
        """Get a command by name, loading export subcommands on first access."""
        self._ensure_export_loaded()
        return super().get_command(ctx, cmd_name)

    def _ensure_export_loaded(self):
        if not self._export_loaded:
            LazyGroup._export_loaded = True
            _register_export_commands()

get_command(ctx, cmd_name)

Get a command by name, loading export subcommands on first access.

Source code in sleap_nn/cli.py

def get_command(self, ctx, cmd_name):
    """Get a command by name, loading export subcommands on first access."""
    self._ensure_export_loaded()
    return super().get_command(ctx, cmd_name)

list_commands(ctx)

List all commands, loading export subcommands on first access.

Source code in sleap_nn/cli.py

def list_commands(self, ctx):
    """List all commands, loading export subcommands on first access."""
    self._ensure_export_loaded()
    return super().list_commands(ctx)

TrainCommand

Bases: Command

Custom command class that overrides help behavior for train command.

Methods:

Name	Description
format_help	Override the help formatting to show custom training help.

Source code in sleap_nn/cli.py

class TrainCommand(Command):
    """Custom command class that overrides help behavior for train command."""

    def format_help(self, ctx, formatter):
        """Override the help formatting to show custom training help."""
        show_training_help()

format_help(ctx, formatter)

Override the help formatting to show custom training help.

Source code in sleap_nn/cli.py

def format_help(self, ctx, formatter):
    """Override the help formatting to show custom training help."""
    show_training_help()

cli()

SLEAP-NN: Neural network backend for training and inference for animal pose estimation.

Use subcommands to run different workflows:

train - Run training workflow (auto-handles multi-GPU) predict - Run inference workflow (new pipeline) track - Run inference/tracking workflow (legacy pipeline) eval - Run evaluation workflow system - Display system information and GPU status

Source code in sleap_nn/cli.py

@click.group(cls=LazyGroup, context_settings=CONTEXT_SETTINGS)
@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 (auto-handles multi-GPU)
    predict  - Run inference workflow (new pipeline)
    track    - Run inference/tracking workflow (legacy pipeline)
    eval     - Run evaluation workflow
    system   - Display system information and GPU status
    """
    pass

config(slp_path, output, auto, pipeline, show_yaml)

Generate training configuration for a SLEAP file.

[Experimental] This feature is experimental and may change in future releases.

Launch an interactive TUI to configure training, or use --auto to generate a config with smart defaults based on your data.

Examples:

Interactive TUI

sleap-nn config labels.slp

Auto-generate config

sleap-nn config labels.slp --auto -o config.yaml

Auto-generate with overrides

sleap-nn config labels.slp --auto --pipeline bottomup

Box detector + box-scaled centered-instance pair

sleap-nn config labels.slp --auto --pipeline topdown_bbox_centernet -o cfg.yaml

Source code in sleap_nn/cli.py

@cli.command(context_settings=CONTEXT_SETTINGS)
@click.argument("slp_path", type=str, required=False, default=None)
@click.option(
    "--output",
    "-o",
    type=str,
    default=None,
    help="Output path for the generated config file(s).",
)
@click.option(
    "--auto",
    is_flag=True,
    default=False,
    help="Auto-generate config without interactive TUI.",
)
@click.option(
    "--pipeline",
    type=click.Choice(
        [
            "single_instance",
            "centroid",
            "bottomup",
            "topdown",
            "multi_class_bottomup",
            "multi_class_topdown",
            "bbox_centernet",
            "bbox_ssd",
            "topdown_bbox_centernet",
            "topdown_bbox_ssd",
        ]
    ),
    default=None,
    help=(
        "Override model pipeline type. 'centroid' generates a single STANDALONE "
        "centroid config (one point per animal); 'topdown' generates paired "
        "centroid + centered_instance configs; 'bbox_centernet' / 'bbox_ssd' "
        "generate a bounding-box detector trained on boxes derived from the poses; "
        "'topdown_bbox_centernet' / 'topdown_bbox_ssd' generate paired detector + "
        "centered_instance configs whose crops are the detected boxes "
        "(crop_mode: bbox)."
    ),
)
@click.option(
    "--show-yaml",
    is_flag=True,
    default=False,
    help="Print generated YAML to stdout.",
)
def config(
    slp_path,
    output,
    auto,
    pipeline,
    show_yaml,
):
    """Generate training configuration for a SLEAP file.

    **[Experimental]** This feature is experimental and may change in future releases.

    Launch an interactive TUI to configure training, or use --auto to
    generate a config with smart defaults based on your data.

    Examples:
        # Interactive TUI
        sleap-nn config labels.slp

        # Auto-generate config
        sleap-nn config labels.slp --auto -o config.yaml

        # Auto-generate with overrides
        sleap-nn config labels.slp --auto --pipeline bottomup

        # Box detector + box-scaled centered-instance pair
        sleap-nn config labels.slp --auto --pipeline topdown_bbox_centernet -o cfg.yaml
    """
    from sleap_nn.config_generator.generator import ConfigGenerator

    # Auto mode (non-interactive)
    if auto:
        if not slp_path:
            click.echo("Error: SLP_PATH is required for --auto mode", err=True)
            raise SystemExit(1)

        gen = ConfigGenerator.from_slp(slp_path)
        gen.auto()

        # Apply overrides
        if pipeline:
            # 'topdown' is a CLI alias for the centroid STAGE (is_topdown=True),
            # which triggers paired centroid + centered_instance generation.
            # 'centroid' is the STANDALONE single-config centroid model
            # (is_topdown=False), emitted via the "centroid_only" pipeline.
            if pipeline == "topdown":
                gen.pipeline("centroid")
            elif pipeline == "centroid":
                gen.pipeline("centroid_only")
            else:
                gen.pipeline(pipeline)

        if show_yaml:
            click.echo(gen.to_yaml())
        elif output:
            gen.save(output)
            if gen.is_topdown:
                # "centroid" (or "detector" for topdown_bbox_*), then "instance".
                for stage, stage_path in gen.topdown_paths(output):
                    click.echo(f"Saved {stage} config to: {stage_path}")
            else:
                click.echo(f"Saved config to: {output}")
        else:
            # Default output path
            slp_stem = Path(slp_path).stem
            output_path = Path(slp_path).parent / f"{slp_stem}_config.yaml"
            gen.save(str(output_path))
            if gen.is_topdown:
                for stage, stage_path in gen.topdown_paths(str(output_path)):
                    click.echo(f"Saved {stage} config to: {stage_path}")
            else:
                click.echo(f"Saved config to: {output_path}")
        return

    # Interactive TUI mode
    try:
        from sleap_nn.config_generator.tui.app import launch_tui

        launch_tui(slp_path)
    except ImportError as e:
        click.echo(f"Error: TUI dependencies not available: {e}", err=True)
        click.echo("Install with: pip install textual", err=True)
        raise SystemExit(1)

eval(**kwargs)

Run evaluation workflow.

Source code in sleap_nn/cli.py

@cli.command(context_settings=CONTEXT_SETTINGS)
@click.option(
    "--ground_truth_path",
    "-g",
    type=str,
    required=True,
    help="Path to ground truth labels file (.slp)",
)
@click.option(
    "--predicted_path",
    "-p",
    type=str,
    required=True,
    help="Path to predicted labels file (.slp)",
)
@click.option("--save_metrics", "-s", type=str, help="Path to save metrics (.npz file)")
@click.option(
    "--oks_stddev",
    type=float,
    default=0.025,
    help="Standard deviation for OKS calculation",
)
@click.option("--oks_scale", type=float, help="Scale factor for OKS calculation")
@click.option(
    "--match_threshold", type=float, default=0.0, help="Threshold for instance matching"
)
@click.option(
    "--user_labels_only/--no-user_labels_only",
    default=True,
    help="Only evaluate user-labeled frames (default: True)",
)
@click.option(
    "--match_method",
    type=click.Choice(["oks", "centroid", "mask", "semantic", "bbox", "auto"]),
    default="auto",
    help=(
        "Matching method: 'oks' (full-skeleton), 'centroid' (single-point "
        "pixel-distance), 'mask' (instance-segmentation mask IoU), 'semantic' "
        "(whole-frame foreground IoU/clDice, no matching), 'bbox' (COCO box AP of "
        "predicted bounding boxes), or 'auto' (bbox when the predictions are "
        "boxes only, centroid when the prediction skeleton is single-node). "
        "Default: auto."
    ),
)
@click.option(
    "--bbox_source",
    type=click.Choice(["auto", "bboxes", "instances", "masks"]),
    default="auto",
    help=(
        "bbox mode: where ground-truth boxes come from (user boxes, or derived "
        "from poses / masks). Pass the model's data_config.bbox_source."
    ),
)
@click.option(
    "--bbox_padding",
    type=float,
    default=0.0,
    help=(
        "bbox mode: pixels of padding on ground-truth boxes derived from poses or "
        "masks (when the ground truth has no boxes). Pass the model's "
        "data_config.bbox_padding."
    ),
)
@click.option(
    "--anchor_part",
    type=str,
    default=None,
    help=(
        "GT skeleton node used to compute ground-truth centroids in centroid "
        "mode. Defaults to the mean of visible nodes when absent (#586)."
    ),
)
@click.option(
    "--centroid_method",
    type=click.Choice(["center_of_mass", "bbox_center", "geometric_median", "anchor"]),
    default=None,
    help=(
        "How ground-truth centroids are derived in centroid mode. Defaults to "
        "the anchor node when --anchor_part is given, else the mean of visible "
        "nodes. Pass the value the model was TRAINED with (its "
        "head_configs.centroid.confmaps.centroid_method) so the metric compares "
        "like with like (#586)."
    ),
)
@click.option(
    "--centroid_fallback",
    type=click.Choice(["center_of_mass", "bbox_center", "geometric_median"]),
    default=None,
    help=(
        "Reduce method used when the --anchor_part node is not visible. "
        "Default: center_of_mass."
    ),
)
@click.option(
    "--count_unmatched_as_fp/--no-count_unmatched_as_fp",
    default=False,
    help=(
        "Count unmatched predictions as false positives in the oks_voc metrics. "
        "Default off, which keeps the numbers earlier versions reported (those "
        "ignore unmatched predictions, so AP is overstated for models that "
        "produce duplicate or spurious detections)."
    ),
)
def eval(**kwargs):
    """Run evaluation workflow."""
    from sleap_nn.evaluation import run_evaluation

    run_evaluation(**kwargs)

eval_tracking(**kwargs)

Evaluate identity persistence of a tracked prediction.

Scores whether tracks keep the right identity across frames -- ID switches, IDF1, MT/PT/ML, fragmentation and track purity -- against tracked ground truth. Complements sleap-nn eval, which scores detection and localization but says nothing about identity.

Also reports TrackEval-faithful CLEAR-MOT, IDF1 and HOTA (the "mot" block of the saved JSON), numbers comparable to published MOT results. Where both blocks name the same metric they can differ: the legacy identity block uses sleap-nn's own definitions.

Both files must be tracked: ground truth needs track set on the detections to score, and the prediction needs tracks from sleap-nn track or sleap-nn predict -t. Exits non-zero when nothing is scoreable (no tracked prediction, no frame in common, or no tracked ground truth).

Examples:

sleap-nn eval-tracking -g gt.slp -p tracked.slp

sleap-nn eval-tracking -g gt.slp -p tracked.slp --carrier mask -s ids.json

Source code in sleap_nn/cli.py

@cli.command("eval-tracking", context_settings=CONTEXT_SETTINGS)
@click.option(
    "--ground_truth_path",
    "-g",
    type=str,
    required=True,
    help="Path to tracked ground truth labels file (.slp)",
)
@click.option(
    "--predicted_path",
    "-p",
    type=str,
    required=True,
    help="Path to tracked predicted labels file (.slp)",
)
@click.option(
    "--save_metrics", "-s", type=str, help="Path to save metrics (.json file)"
)
@click.option(
    "--carrier",
    type=click.Choice(["auto", "pose", "mask"]),
    default="auto",
    help=(
        "What carries identity: 'pose' (instances, matched by OKS, or by "
        "centroid distance for single-node skeletons), 'mask' (segmentation "
        "masks, matched by IoU), or 'auto' (the carrier the prediction's tracks "
        "are on -- the one with more tracked detections, ties to pose). A "
        "detection without a track takes the track of its linked instance or "
        "mask. Default: auto."
    ),
)
@click.option(
    "--match_threshold",
    type=float,
    default=None,
    help=(
        "Matching threshold for a predicted detection to count as matched to a "
        "ground-truth one: minimum OKS (pose) or mask IoU (mask), or -- for "
        "single-node (centroid) skeletons, which match by distance -- the "
        "maximum centroid distance in PIXELS. Default: 0.5 for OKS/IoU, 50 px "
        "for distance (as in `sleap-nn eval --match_method centroid`)."
    ),
)
@click.option(
    "--mt_threshold",
    type=float,
    default=0.8,
    help=(
        "Coverage at or above which a GT trajectory is mostly-tracked (MT). "
        "Legacy identity block only; the TrackEval `mot` block always uses its "
        "fixed MT > 0.8."
    ),
)
@click.option(
    "--ml_threshold",
    type=float,
    default=0.2,
    help=(
        "Coverage below which a GT trajectory is mostly-lost (ML). Legacy "
        "identity block only; the TrackEval `mot` block always uses its fixed "
        "ML < 0.2."
    ),
)
@click.option(
    "--user_labels_only/--no-user_labels_only",
    default=False,
    help=(
        "Drop predicted detections from the GROUND-TRUTH side (default: False, "
        "unlike `sleap-nn eval`). Tracked ground truth is usually predicted "
        "poses with tracks assigned afterwards, so filtering by type would "
        "discard it. Pass this only for user-labeled GT that also carries "
        "stale predictions from an earlier run."
    ),
)
@click.option(
    "--global_identity/--per_video_identity",
    default=False,
    help=(
        "Score track names as GLOBAL identities shared across videos (names "
        "from a multi-class ID model, or a project proofread with one name per "
        "animal): a trajectory then spans every video, and keeping an animal's "
        "name from one video to the next is part of the score. Default: per "
        "video -- each (video, track name) is its own trajectory, since "
        "trackers name tracks per video."
    ),
)
@click.option(
    "--oks_stddev",
    type=click.FloatRange(min=0, min_open=True),
    default=0.025,
    show_default=True,
    callback=_require_finite,
    help=(
        "Per-keypoint spread of the OKS used to match poses (both the identity "
        "and the MOT metrics); a finite number > 0. Larger values match looser "
        "poses. Ignored for masks and single-node skeletons."
    ),
)
def eval_tracking(**kwargs):
    """Evaluate identity persistence of a tracked prediction.

    Scores whether tracks keep the right identity across frames -- ID switches,
    IDF1, MT/PT/ML, fragmentation and track purity -- against tracked ground
    truth. Complements `sleap-nn eval`, which scores detection and localization
    but says nothing about identity.

    Also reports TrackEval-faithful CLEAR-MOT, IDF1 and HOTA (the "mot" block of
    the saved JSON), numbers comparable to published MOT results. Where both
    blocks name the same metric they can differ: the legacy identity block uses
    sleap-nn's own definitions.

    Both files must be tracked: ground truth needs `track` set on the detections
    to score, and the prediction needs tracks from `sleap-nn track` or
    `sleap-nn predict -t`. Exits non-zero when nothing is scoreable (no tracked
    prediction, no frame in common, or no tracked ground truth).

    Examples:
        sleap-nn eval-tracking -g gt.slp -p tracked.slp

        sleap-nn eval-tracking -g gt.slp -p tracked.slp --carrier mask -s ids.json
    """
    from sleap_nn.evaluation import NothingToScoreError, run_identity_evaluation

    try:
        run_identity_evaluation(**kwargs, raise_if_empty=True)
    except NothingToScoreError as e:
        raise click.ClickException(str(e)) from e

infer(**kwargs)

Deprecated alias for predict. Use sleap-nn predict instead.

Source code in sleap_nn/cli.py

@cli.command("infer", hidden=True, context_settings=CONTEXT_SETTINGS)
@_common_inference_options
def infer(**kwargs):
    """Deprecated alias for ``predict``. Use ``sleap-nn predict`` instead."""
    import warnings

    warnings.warn(
        "sleap-nn infer is deprecated. Use sleap-nn predict instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    return _run_inference_impl(**kwargs)

info(path)

Display model configuration and evaluation metrics.

PATH can be a trained model directory or a training config YAML file.

If a model directory is given, shows config summary, training results, evaluation metrics (if available), and files in the directory.

If a config YAML is given, shows only the config summary.

Examples:

sleap-nn info path/to/model_dir

sleap-nn info path/to/training_config.yaml

Source code in sleap_nn/cli.py

@cli.command(context_settings=CONTEXT_SETTINGS)
@click.argument("path", type=str)
def info(path):
    """Display model configuration and evaluation metrics.

    PATH can be a trained model directory or a training config YAML file.

    If a model directory is given, shows config summary, training results,
    evaluation metrics (if available), and files in the directory.

    If a config YAML is given, shows only the config summary.

    Examples:
        sleap-nn info path/to/model_dir

        sleap-nn info path/to/training_config.yaml
    """
    from sleap_nn.model_info import print_model_info

    print_model_info(path)

is_config_path(arg)

Check if an argument looks like a config file path.

Returns True if the arg ends with .yaml or .yml.

Source code in sleap_nn/cli.py

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")

parse_path_map(ctx, param, value)

Parse (old, new) path pairs into a dictionary for path mapping options.

Source code in sleap_nn/cli.py

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

predict(**kwargs)

Run inference on videos or labels files.

Single unified inference entry point.

Source code in sleap_nn/cli.py

@cli.command(context_settings=CONTEXT_SETTINGS)
@_common_inference_options
def predict(**kwargs):
    """Run inference on videos or labels files.

    Single unified inference entry point.
    """
    return _run_inference_impl(**kwargs)

print_version(ctx, param, value)

Print version and exit.

Source code in sleap_nn/cli.py

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()

show_training_help()

Display training help information with rich formatting.

Source code in sleap_nn/cli.py

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
```

**Train on a different labels file** (replaces the config's labels paths):
```bash
sleap-nn train config.yaml --labels labels.slp
sleap-nn train config.yaml -l train.slp --val-labels val.slp
```

**Check a config without training (nothing is written to disk):**
```bash
sleap-nn train config.yaml --dry-run
```

**Resume training** (writes to a new `<run_name>-1` folder next to the original):
```bash
sleap-nn train --resume path/to/run_dir
sleap-nn train --resume path/to/run_dir trainer_config.max_epochs=200
sleap-nn train --resume path/to/run_dir/best.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/
"""
    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",
        )
    )

split_config_path(config_path)

Split a full config path into (config_dir, config_name).

Parameters:

Name	Type	Description	Default
config_path	str	Full path to a config file.	required

Returns:

Type	Description
tuple	Tuple of (config_dir, config_name) where config_dir is an absolute path.

Source code in sleap_nn/cli.py

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

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.

Source code in sleap_nn/cli.py

@cli.command(context_settings=CONTEXT_SETTINGS)
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()

track(**kwargs)

Run Inference and Tracking workflow (legacy pipeline).

This command uses the legacy run_inference pipeline. For the new inference pipeline, use sleap-nn predict.

Source code in sleap_nn/cli.py

@cli.command(context_settings=CONTEXT_SETTINGS)
@click.option(
    "--data_path",
    "-i",
    type=str,
    required=True,
    help="Path to data to predict on. This can be a labels (.slp) file or any supported video format.",
)
@click.option(
    "--model_paths",
    "-m",
    multiple=True,
    help="Path to a trained model directory, or to its best.ckpt or training_config.yaml/.json file (all resolve to the model directory). Multiple models can be specified, each preceded by --model_paths.",
)
@click.option(
    "--output_path",
    "-o",
    type=str,
    default=None,
    help="The output filename to use for the predicted data. If not provided, defaults to '[data_path].slp'.",
)
@click.option(
    "--device",
    "-d",
    type=str,
    default="auto",
    help="Device on which torch.Tensor will be allocated. One of the ('cpu', 'cuda', 'mps', 'auto'). Default: 'auto' (based on available backend either cuda, mps or cpu is chosen). If `cuda` is available, you could also use `cuda:0` to specify the device.",
)
@click.option(
    "--batch_size",
    "-b",
    type=int,
    default=4,
    help="Number of frames to predict at a time. Larger values result in faster inference speeds, but require more memory.",
)
@click.option(
    "--tracking",
    "-t",
    is_flag=True,
    default=False,
    help="If True, runs tracking on the predicted instances.",
)
@click.option(
    "-n",
    "--max_instances",
    type=int,
    default=None,
    help="Limit maximum number of instances in multi-instance models. Not available for ID models. Defaults to None.",
)
@click.option(
    "--backbone_ckpt_path",
    type=str,
    default=None,
    help="To run inference on any `.ckpt` other than `best.ckpt` from the `model_paths` dir, the path to the `.ckpt` file should be passed here.",
)
@click.option(
    "--head_ckpt_path",
    type=str,
    default=None,
    help="Path to `.ckpt` file if a different set of head layer weights are to be used. If `None`, the `best.ckpt` from `model_paths` dir is used (or the ckpt from `backbone_ckpt_path` if provided.)",
)
@click.option(
    "--max_height",
    type=int,
    default=None,
    help="Maximum height the image should be padded to. If not provided, the values from the training config are used. Default: None.",
)
@click.option(
    "--max_width",
    type=int,
    default=None,
    help="Maximum width the image should be padded to. If not provided, the values from the training config are used. Default: None.",
)
@click.option(
    "--input_scale",
    type=float,
    default=None,
    help="Scale factor to apply to the input image. If not provided, the values from the training config are used. Default: None.",
)
@click.option(
    "--ensure_rgb/--no-ensure_rgb",
    default=None,
    help="If True, input images will have 3 channels (RGB). Single-channel images are replicated along the channel axis. If False, RGB conversion is disabled. If not provided, the values from the training config are used. Default: None.",
)
@click.option(
    "--ensure_grayscale/--no-ensure_grayscale",
    default=None,
    help="If True, input images will be converted to single-channel grayscale. If False, grayscale conversion is disabled. If not provided, the values from the training config are used. Default: None.",
)
@click.option(
    "--anchor_part",
    type=str,
    default=None,
    help="The node name to use as the anchor for the centroid. If not provided, the anchor part in the `training_config.yaml` is used. Default: `None`.",
)
@click.option(
    "--only_labeled_frames",
    is_flag=True,
    default=False,
    help="Only run inference on user labeled frames when running on labels dataset. This is useful for generating predictions to compare against ground truth.",
)
@click.option(
    "--only_suggested_frames",
    is_flag=True,
    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,
    default=False,
    help=("Clear empty frames that did not have predictions before saving to output."),
)
@click.option(
    "--video_index",
    type=int,
    default=None,
    help="Integer index of video in .slp file to predict on. To be used with an .slp path as an alternative to specifying the video path: predicts on every frame of that video (or --frames, as source frame indices), and --exclude_user_labeled / --only_* filter those frames using the .slp's annotations. A video embedded in a .pkg.slp is scoped to its labeled frames instead.",
)
@click.option(
    "--video_dataset", type=str, default=None, help="The dataset for HDF5 videos."
)
@click.option(
    "--video_input_format",
    type=str,
    default="channels_last",
    help="The input_format for HDF5 videos.",
)
@click.option(
    "--frames",
    type=str,
    default="",
    help="List of frames to predict when running on a video. Can be specified as a comma separated list (e.g. 1,2,3) or a range separated by hyphen (e.g., 1-3, for 1,2,3). If not provided, defaults to predicting on the entire video.",
)
@click.option(
    "--integral_patch_size",
    type=int,
    default=5,
    help="Size of patches to crop around each rough peak as an integer scalar. Default: 5.",
)
@click.option(
    "--max_edge_length_ratio",
    type=float,
    default=0.25,
    help="The maximum expected length of a connected pair of points as a fraction of the image size. Candidate connections longer than this length will be penalized during matching. Default: 0.25.",
)
@click.option(
    "--dist_penalty_weight",
    type=float,
    default=1.0,
    help="A coefficient to scale weight of the distance penalty as a scalar float. Set to values greater than 1.0 to enforce the distance penalty more strictly. Default: 1.0.",
)
@click.option(
    "--n_points",
    type=int,
    default=10,
    help="Number of points to sample along the line integral. Default: 10.",
)
@click.option(
    "--min_instance_peaks",
    type=float,
    default=0,
    help="Minimum number of peaks the instance should have to be considered a real instance. Instances with fewer peaks than this will be discarded (useful for filtering spurious detections). Default: 0.",
)
@click.option(
    "--min_line_scores",
    type=float,
    default=0.25,
    help="Minimum line score (between -1 and 1) required to form a match between candidate point pairs. Useful for rejecting spurious detections when there are no better ones. Default: 0.25.",
)
@click.option(
    "--queue_maxsize",
    type=int,
    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. If `input_scale` is provided, then the cropped image will be resized according to `input_scale`.",
)
@click.option(
    "--peak_threshold",
    type=float,
    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(
    "--filter_min_visible_nodes",
    type=int,
    default=0,
    help=(
        "Minimum number of visible (non-NaN) keypoints required. "
        "Instances with fewer visible nodes are removed. (default: 0, no filtering)"
    ),
)
@click.option(
    "--filter_min_visible_node_fraction",
    type=float,
    default=0.0,
    help=(
        "Minimum fraction of skeleton nodes that must be visible. "
        "Value should be in [0, 1]. For example, 0.5 requires at least half "
        "of the skeleton's nodes to be detected. (default: 0.0, no filtering)"
    ),
)
@click.option(
    "--filter_min_mean_node_score",
    type=float,
    default=0.0,
    help=(
        "Minimum mean confidence score across visible nodes. "
        "Instances with lower mean node scores are removed. (default: 0.0, no filtering)"
    ),
)
@click.option(
    "--filter_min_instance_score",
    type=float,
    default=0.0,
    help=(
        "Minimum overall instance confidence score. "
        "Instances with lower scores are removed. (default: 0.0, no filtering)"
    ),
)
@click.option(
    "--integral_refinement",
    type=str,
    default="integral",
    help="Sub-pixel peak refinement: 'integral' (integral regression), 'local' (quarter-pixel offset towards the higher neighbor) or 'none' (grid-aligned peaks). Default: 'integral'.",
)
@click.option(
    "--tracking_window_size",
    type=int,
    default=5,
    help="Number of frames to look for in the candidate instances to match with the current detections.",
)
@click.option(
    "--min_new_track_points",
    type=int,
    default=0,
    help="We won't spawn a new track for an instance with fewer than this many points.",
)
@click.option(
    "--candidates_method",
    type=str,
    default="fixed_window",
    help="Either of `fixed_window` or `local_queues`. In fixed window method, candidates from the last `window_size` frames. In local queues, last `window_size` instances for each track ID is considered for matching against the current detection.",
)
@click.option(
    "--min_match_points",
    type=int,
    default=0,
    help="Minimum non-NaN points for match candidates.",
)
@click.option(
    "--features",
    type=str,
    default="keypoints",
    help="Feature representation for the candidates to update current detections. One of [`keypoints`, `centroids`, `bboxes`, `masks`, `embeddings`]. `embeddings` tracks by the appearance vector attached by the embedding model (pair with `--scoring_method cosine_sim`).",
)
@click.option(
    "--scoring_method",
    type=str,
    default="oks",
    help="Method to compute association score between features from the current frame and the previous tracks. One of [`oks`, `cosine_sim`, `iou`, `mask_iou`, `euclidean_dist`]. `cosine_sim` pairs with `--features embeddings`.",
)
@click.option(
    "--scoring_reduction",
    type=str,
    default="mean",
    help="Method to aggregate and reduce multiple scores if there are several detections associated with the same track. One of [`mean`, `max`, `robust_quantile`].",
)
@click.option(
    "--robust_best_instance",
    type=float,
    default=1.0,
    help="If the value is between 0 and 1 (excluded), use a robust quantile similarity score for the track. If the value is 1, use the max similarity (non-robust). For selecting a robust score, 0.95 is a good value.",
)
@click.option(
    "--track_matching_method",
    type=str,
    default="hungarian",
    help="Track matching algorithm. One of `hungarian`, `greedy`.",
)
@click.option(
    "--max_tracks",
    type=int,
    default=None,
    help="Maximum number of new tracks to be created to avoid redundant tracks. (only for local queues candidate)",
)
@click.option(
    "--use_flow",
    is_flag=True,
    default=False,
    help="If True, `FlowShiftTracker` is used, where the poses are matched using optical flow shifts.",
)
@click.option(
    "--of_img_scale",
    type=float,
    default=1.0,
    help="Factor to scale the images by when computing optical flow. Decrease this to increase performance at the cost of finer accuracy. Sometimes decreasing the image scale can improve performance with fast movements.",
)
@click.option(
    "--of_window_size",
    type=int,
    default=21,
    help="Optical flow window size to consider at each pyramid scale level.",
)
@click.option(
    "--of_max_levels",
    type=int,
    default=3,
    help="Number of pyramid scale levels to consider. This is different from the scale parameter, which determines the initial image scaling.",
)
@click.option(
    "--use_kalman",
    is_flag=True,
    default=False,
    help="If True, `KalmanShiftTracker` is used: poses are predicted with a per-track constant-velocity Kalman filter. Requires --tracking_target_instance_count (or --max_tracks/--max_instances) and is mutually exclusive with --use_flow.",
)
@click.option(
    "--kf_track_features",
    type=click.Choice(["centroid", "keypoints"]),
    default="centroid",
    help="What the Kalman motion model tracks: 'centroid' (default; rigid, stable) or 'keypoints' (per-node poses; noisier — pair with --oks_stddev or --features bboxes --scoring_method iou). (only if --use_kalman)",
)
@click.option(
    "--oks_stddev",
    type=float,
    default=None,
    help="OKS keypoint-spread normalization constant for `oks` scoring. Larger is more tolerant of localization error (useful with --kf_track_features keypoints). Default: 0.025.",
)
@click.option(
    "--kf_init_frame_count",
    type=int,
    default=10,
    help="Number of warm-up frames tracked with the base path before the Kalman filters are fit via EM. (only if --use_kalman)",
)
@click.option(
    "--kf_node_indices",
    type=str,
    default=None,
    callback=_parse_int_list,
    help="Comma-separated skeleton node indices to track with the motion model (e.g. '0,1,2'). Empty/unset uses all nodes. (only if --use_kalman)",
)
@click.option(
    "--kf_reset_gap_size",
    type=int,
    default=5,
    help="Number of consecutive missed frames after which a stale track's Kalman filter is reset. (only if --use_kalman)",
)
@click.option(
    "--post_connect_single_breaks",
    is_flag=True,
    default=False,
    help="If True and `max_tracks` is not None with local queues candidate method, connects track breaks when exactly one track is lost and exactly one new track is spawned in the frame.",
)
@click.option(
    "--tracking_target_instance_count",
    type=int,
    default=None,
    help="Target number of instances to track per frame. (default: 0)",
)
@click.option(
    "--tracking_pre_cull_to_target",
    type=int,
    default=0,
    help=(
        "If non-zero and target_instance_count is also non-zero, then cull instances "
        "over target count per frame *before* tracking. (default: 0)"
    ),
)
@click.option(
    "--tracking_pre_cull_iou_threshold",
    type=float,
    default=0,
    help=(
        "If non-zero and pre_cull_to_target also set, then use IOU threshold to remove "
        "overlapping instances over count *before* tracking. (default: 0)"
    ),
)
@click.option(
    "--tracking_clean_instance_count",
    type=int,
    default=0,
    help="Target number of instances to clean *after* tracking. (default: 0)",
)
@click.option(
    "--tracking_clean_iou_threshold",
    type=float,
    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 (legacy pipeline).

    This command uses the legacy ``run_inference`` pipeline. For the new
    inference pipeline, use ``sleap-nn predict``.
    """
    from sleap_nn import redirect_logs_to_stderr
    from sleap_nn.legacy_predict import frame_list, run_inference

    if kwargs.get("gui"):
        redirect_logs_to_stderr()

    if "model_paths" in kwargs and kwargs["model_paths"]:
        kwargs["model_paths"] = list(kwargs["model_paths"])
    else:
        kwargs["model_paths"] = None

    if "frames" in kwargs and kwargs["frames"]:
        kwargs["frames"] = frame_list(kwargs["frames"])
    else:
        kwargs["frames"] = None

    # Pop new-pipeline-only flags that track doesn't use. NOTE: `gui` is NOT
    # popped — legacy run_inference accepts+honors it (predict.py sets
    # predictor.gui, switching to JSON-per-line progress); popping it silently
    # dropped `track --gui` GUI integration (#583).
    kwargs.pop("paf_workers", None)
    kwargs.pop("cpu_workers", None)
    kwargs.pop("stream_to_file", None)
    kwargs.pop("write_interval", None)
    # `--runtime` selects the exported-model runtime for `sleap-nn predict`;
    # `track` only runs trained checkpoints, so it is inert here.
    kwargs.pop("runtime", None)
    # New-flow-only flags, inert here. `track` cannot do centroid-only
    # inference; a lone centroid model is rejected downstream by run_inference
    # with a message pointing to `sleap-nn predict`.
    kwargs.pop("centroid_only", None)
    kwargs.pop("centroid_peak_threshold", None)
    kwargs.pop("centroid_output", None)
    kwargs.pop("filter_min_centroid_distance", None)
    # SAM prompted-mask flags are new-flow-only (legacy run_inference rejects
    # them); inert here.
    kwargs.pop("mask_backend", None)
    kwargs.pop("sam_checkpoint", None)
    kwargs.pop("sam_model_type", None)
    kwargs.pop("sam_prompt_mode", None)
    kwargs.pop("sam_anchor_ind", None)
    kwargs.pop("sam_disjointify_masks", None)
    kwargs.pop("sam3_model_id", None)
    kwargs.pop("overlay_path", None)

    return run_inference(**kwargs)

train(config, config_name, config_dir, video_paths, video_path_map, prefix_map, video_config, labels, val_labels, resume, dry_run, 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.yaml --labels path/to/labels.slp sleap-nn train --resume path/to/run_dir trainer_config.max_epochs=200 sleap-nn train config.yaml --dry-run

Source code in sleap_nn/cli.py

@cli.command(cls=TrainCommand, context_settings=CONTEXT_SETTINGS)
@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)"
)
@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.option(
    "--labels",
    "-l",
    multiple=True,
    metavar="SLP",
    help="Training labels (.slp), replacing data_config.train_labels_path. "
    "Can be specified multiple times.",
)
@click.option(
    "--val-labels",
    multiple=True,
    metavar="SLP",
    help="Validation labels (.slp), replacing data_config.val_labels_path. "
    "Without it, the config's validation labels are used, or a "
    "validation_fraction split of the training labels if it has none. "
    "Can be specified multiple times.",
)
@click.option(
    "--resume",
    type=str,
    default=None,
    metavar="RUN_DIR_OR_CKPT",
    help="Resume a previous run: a run folder (uses its last.ckpt, else best.ckpt) "
    "or a .ckpt file. Uses that run's training_config.yaml (overrides still "
    "apply) and writes to a new folder <run_name>-1 next to it, leaving the "
    "original run untouched.",
)
@click.option(
    "--dry-run",
    is_flag=True,
    default=False,
    help="Check the config and data without training: run all setup plus one "
    "training and one validation batch on a single device, with the config's "
    "settings on one batch of frames; check the run can write its configs and a "
    "checkpoint next to the run folder (removed afterwards); and print a summary "
    "of the run.",
)
@click.argument("overrides", nargs=-1, type=click.UNPROCESSED)
def train(
    config,
    config_name,
    config_dir,
    video_paths,
    video_path_map,
    prefix_map,
    video_config,
    labels,
    val_labels,
    resume,
    dry_run,
    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.yaml --labels path/to/labels.slp
        sleap-nn train --resume path/to/run_dir trainer_config.max_epochs=200
        sleap-nn train config.yaml --dry-run
    """
    import hydra
    import sleap_io as sio
    from omegaconf import OmegaConf
    from sleap_nn.train import run_training

    # 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
    resume_target = None
    if resume is not None:
        if config_from_positional or config or config_name:
            raise click.UsageError(
                "--resume uses the run's own training_config.yaml; don't pass a "
                "config too. Use overrides (e.g. trainer_config.max_epochs=200) to "
                "change settings."
            )
        resume_target = _resolve_resume_target(resume)
        config_dir, config_name = (
            resume_target[0].as_posix(),
            "training_config.yaml",
        )
    elif 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

    labels = _resolve_labels_option(
        labels, "--labels", "data_config.train_labels_path", overrides
    )
    val_labels = _resolve_labels_option(
        val_labels, "--val-labels", "data_config.val_labels_path", overrides
    )

    # 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):
        cfg = hydra.compose(config_name=config_name, overrides=overrides)
        cfg = _adjust_train_config(
            cfg, overrides, resume_target, dry_run, labels, val_labels
        )
        # Logged here, not in `_adjust_train_config`: every route below composes
        # the config again.
        if resume_target is not None:
            logger.info(
                f"Resuming from {cfg.trainer_config.resume_ckpt_path} into new run "
                f"folder {(Path(cfg.trainer_config.ckpt_dir) / cfg.trainer_config.run_name).as_posix()}"
            )
        config_val_labels = OmegaConf.select(cfg, "data_config.val_labels_path")
        if labels and not val_labels and config_val_labels:
            logger.info(
                f"Validating on the config's val_labels_path ({list(config_val_labels)}); "
                "pass --val-labels to use different validation labels."
            )

        # Validate config
        if not hasattr(cfg, "model_config") or not cfg.model_config:
            click.echo(
                "No model config found! Use `sleap-nn train --help` for more information."
            )
            raise click.Abort()

        num_devices = _get_num_devices_from_config(cfg)

        # Check if run_name is already set (for synchronization across DDP ranks)
        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
    # We need to re-spawn if EITHER:
    # 1. Not in module context (__main__.__spec__ is None) - required for DDP on
    #    Windows/macOS where multiprocessing uses 'spawn' and needs to know what
    #    module to re-import. See: https://github.com/talmolab/sleap/issues/2656
    # 2. run_name is not set - required for synchronization so all DDP ranks use
    #    the same run_name (otherwise each rank generates different timestamps)
    needs_respawn = _needs_module_respawn() or not run_name_is_set
    if num_devices > 1 and needs_respawn:
        logger.info(
            f"Detected {num_devices} devices, re-spawning with module context for DDP..."
        )

        # 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 = _adjust_train_config(
                cfg, overrides, resume_target, dry_run, labels, val_labels
            )
            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)
        cfg = _adjust_train_config(
            cfg, overrides, resume_target, dry_run, labels, val_labels
        )

        logger.info("Input config:")
        logger.info("\n" + OmegaConf.to_yaml(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,
            dry_run=dry_run,
        )