sleap_nn

sleap_nn

Main module for sleap_nn package.

Modules:

Name	Description
architectures	Modules related to model architectures.
cli	Unified CLI for SLEAP-NN using rich-click for styled output.
config	Configuration modules for sleap-nn.
config_generator	Config generator for SLEAP-NN training configurations.
data	Modules related to data loading and processing.
evaluation	This module is to compute evaluation metrics for trained models.
export	Export utilities for sleap-nn.
inference	Inference-related modules.
legacy_models	Utilities for loading legacy SLEAP models.
legacy_predict	Entry point for running inference.
model_info	Model information display for trained models and configs.
system_info	System diagnostics and compatibility checking for sleap-nn.
tracking	Tracker related modules.
train	Entry point for sleap_nn training.
training	Training-related modules.

Functions:

Name	Description
load_metrics	Load metrics from a model folder or metrics file.
load_models	Load trained model(s) into a ready-to-run :class:Predictor.

__dir__()

Include the lazily-exposed names in dir(sleap_nn) for discoverability.

Source code in sleap_nn/__init__.py

def __dir__():
    """Include the lazily-exposed names in ``dir(sleap_nn)`` for discoverability."""
    return sorted(set(globals()) | set(_LAZY_ATTRS))

__getattr__(name)

Resolve lazily-exposed inference entry points (predict, Predictor).

Source code in sleap_nn/__init__.py

def __getattr__(name):
    """Resolve lazily-exposed inference entry points (``predict``, ``Predictor``)."""
    if name in _LAZY_ATTRS:
        from sleap_nn import inference

        return getattr(inference, _LAZY_ATTRS[name])
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

load_metrics(path, split='test', dataset_idx=0)

Load metrics from a model folder or metrics file.

This function supports both the new format (single "metrics" key) and the old format (individual metric keys at top level). It also handles both old and new file naming conventions in model folders.

Parameters:

Name	Type	Description	Default
path	str	Path to a model folder or metrics file (.npz).	required
split	str	Name of the split to load. Must be "train", "val", or "test". Default: "test". If "test" is not found, falls back to "val". Ignored if path points directly to a .npz file.	'test'
dataset_idx	int	Index of the dataset (for multi-dataset training). Default: 0. Ignored if path points directly to a .npz file.	0

Returns:

Type	Description
dict	Dictionary containing metrics with keys: voc_metrics, mOKS, distance_metrics, pck_metrics, visibility_metrics.

Raises:

Type	Description
FileNotFoundError	If no metrics file is found.

Examples:

>>> # Load from model folder (tries test, falls back to val)
>>> metrics = load_metrics("/path/to/model")
>>> print(metrics["mOKS"]["mOKS"])

>>> # Load specific split and dataset
>>> metrics = load_metrics("/path/to/model", split="val", dataset_idx=1)

>>> # Load directly from npz file
>>> metrics = load_metrics("/path/to/metrics.val.0.npz")

Source code in sleap_nn/evaluation.py

def load_metrics(
    path: str,
    split: str = "test",
    dataset_idx: int = 0,
) -> dict:
    """Load metrics from a model folder or metrics file.

    This function supports both the new format (single "metrics" key) and the old
    format (individual metric keys at top level). It also handles both old and new
    file naming conventions in model folders.

    Args:
        path: Path to a model folder or metrics file (.npz).
        split: Name of the split to load. Must be "train", "val", or "test".
            Default: "test". If "test" is not found, falls back to "val".
            Ignored if path points directly to a .npz file.
        dataset_idx: Index of the dataset (for multi-dataset training).
            Default: 0. Ignored if path points directly to a .npz file.

    Returns:
        Dictionary containing metrics with keys: voc_metrics, mOKS,
        distance_metrics, pck_metrics, visibility_metrics.

    Raises:
        FileNotFoundError: If no metrics file is found.

    Examples:
        >>> # Load from model folder (tries test, falls back to val)
        >>> metrics = load_metrics("/path/to/model")
        >>> print(metrics["mOKS"]["mOKS"])

        >>> # Load specific split and dataset
        >>> metrics = load_metrics("/path/to/model", split="val", dataset_idx=1)

        >>> # Load directly from npz file
        >>> metrics = load_metrics("/path/to/metrics.val.0.npz")
    """
    path = Path(path)

    if path.suffix == ".npz":
        metrics_path = path
    else:
        metrics_path = _find_metrics_file(path, split, dataset_idx)

    if not metrics_path.exists():
        raise FileNotFoundError(f"Metrics file not found at {metrics_path}")

    return _load_npz_metrics(metrics_path)

load_models(model_paths, **kwargs)

Load trained model(s) into a ready-to-run :class:Predictor.

A discoverable, top-level convenience wrapper around :meth:sleap_nn.inference.Predictor.from_model_paths. Pass one model directory (single-instance / bottom-up / centroid) or a centroid + centered-instance pair (top-down); a lone centroid directory is auto-detected. Accepts every keyword argument of from_model_paths (e.g. device, batch_size, peak_threshold, tracker_config) and returns a reusable Predictor you can call .predict(...) on repeatedly.

Example

import sleap_nn predictor = sleap_nn.load_models( ... ["models/centroid/", "models/centered_instance/"], device="cuda" ... ) labels = predictor.predict("video.mp4")

For a one-shot call, use :func:sleap_nn.predict instead.

Source code in sleap_nn/__init__.py

def load_models(model_paths, **kwargs):
    """Load trained model(s) into a ready-to-run :class:`Predictor`.

    A discoverable, top-level convenience wrapper around
    :meth:`sleap_nn.inference.Predictor.from_model_paths`. Pass one model
    directory (single-instance / bottom-up / centroid) or a centroid +
    centered-instance pair (top-down); a lone centroid directory is
    auto-detected. Accepts every keyword argument of ``from_model_paths``
    (e.g. ``device``, ``batch_size``, ``peak_threshold``, ``tracker_config``)
    and returns a reusable ``Predictor`` you can call ``.predict(...)`` on
    repeatedly.

    Example:
        >>> import sleap_nn
        >>> predictor = sleap_nn.load_models(
        ...     ["models/centroid/", "models/centered_instance/"], device="cuda"
        ... )
        >>> labels = predictor.predict("video.mp4")

    For a one-shot call, use :func:`sleap_nn.predict` instead.
    """
    from sleap_nn.inference import Predictor

    return Predictor.from_model_paths(model_paths, **kwargs)

redirect_logs_to_stderr()

Redirect sleap_nn's log sink to stderr.

Call this once, early, whenever stdout must be reserved for a machine-readable channel -- e.g. the CLI's --gui mode, which emits one JSON progress line per batch on stdout for a GUI subprocess reader to parse (see sleap_nn.cli._gui_progress_callback). Without this, a plain-text log line (e.g. "Loaded inference model | ...") can interleave with the JSON lines on the same fd and break a naive per-line json.loads() reader. Idempotent -- safe to call more than once.

Source code in sleap_nn/__init__.py

def redirect_logs_to_stderr() -> None:
    """Redirect sleap_nn's log sink to stderr.

    Call this once, early, whenever stdout must be reserved for a
    machine-readable channel -- e.g. the CLI's ``--gui`` mode, which emits one
    JSON progress line per batch on stdout for a GUI subprocess reader to
    parse (see ``sleap_nn.cli._gui_progress_callback``). Without this, a
    plain-text log line (e.g. "Loaded inference model | ...") can interleave
    with the JSON lines on the same fd and break a naive per-line
    ``json.loads()`` reader. Idempotent -- safe to call more than once.
    """
    logger.remove()
    _add_default_sink("stderr")