Skip to content

SLEAP-NN

CI codecov code Release PyPI - Python Version PyPI

SLEAP-NN is the deep learning engine that powers SLEAP (Social LEAP Estimates Animal Poses), providing neural network architectures for multi-instance animal pose estimation and tracking. Built on PyTorch, SLEAP-NN offers an end-to-end training workflow, supporting multiple model types (Single Instance, Top-Down, Bottom-Up, Multi-Class), and seamless integration with SLEAP's GUI and command-line tools.

✨ Features

  • End-to-end workflow: Streamlined training and inference pipelines, including a multi-instance tracking worklow with a flow-shift-based tracker.
  • Multiple model architectures: Support for single-instance, top-down, and bottom-up pose estimation models using a variety of backbones, including highly-customizable UNet, ConvNeXt, and Swin Transformer.
  • PyTorch Lightning integration: Built on PyTorch Lightning for fast and scalable training, with support for multi-GPU and distributed training.
  • Flexible configuration: Hydra and OmegaConf based config system to validate training parameters and enable reproducible experiments.

🚀 Quick Start

Let's start SLEAPiNNg !!! 🐭🐭

Prerequisite: uv installation

Install uv, a fast Python package manager for modern projects: 

# macOS/Linux
curl -LsSf https://astral.sh/uv/install.sh | sh

# Windows
powershell -c "irm https://astral.sh/uv/install.ps1 | iex"

Step - 1 : Set Up Your Configuration

Create a config.yaml file for your experiment. Use a sample config from docs/sample_configs.

Step - 2 : Train a Model

We use uvx here which automatically installs sleap-nn from PyPI with all dependencies and runs the command in a single step. uvx automatically installs sleap-nn and runs your command inside a temporary virtual environment (venv). This means each run is fully isolated and leaves no trace on your system—perfect for trying out sleap-nn without any permanent installation. Check out Installation docs for different installation options.

Quick Start Data: Download sample training data from here and validation data from here for quick experimentation.

uvx --from "sleap-nn[torch]" --index-url https://pypi.org/simple --extra-index-url https://download.pytorch.org/whl/cu128 sleap-nn train --config-name config.yaml --config-dir /path/to/config_dir/ "data_config.train_labels_path=[train.pkg.slp]" "data_config.val_labels_path=[val.pkg.slp]"

uvx --from "sleap-nn[torch]" --index-url https://pypi.org/simple --extra-index-url https://download.pytorch.org/whl/cpu sleap-nn train --config-name config.yaml --config-dir /path/to/config_dir/ "data_config.train_labels_path=[train.pkg.slp]" "data_config.val_labels_path=[val.pkg.slp]"

uvx sleap-nn[torch] train --config-name config.yaml --config-dir /path/to/config_dir/ "data_config.train_labels_path=[train.pkg.slp]" "data_config.val_labels_path=[val.pkg.slp]"

Info

  • For more information on which CUDA version to use for your system, see the PyTorch installation guide.
    The --extra-index-url in the install command should match the CUDA version you need (e.g., https://download.pytorch.org/whl/cuda118 for CUDA 11.8, https://download.pytorch.org/whl/cuda128 for CUDA 12.8, etc.).
  • On macOS, MPS (Metal Performance Shaders) is automatically enabled for Apple Silicon acceleration.

Step - 3 : Run Inference on the Trained Model

To run inference:

uvx --from "sleap-nn[torch]" --index-url https://pypi.org/simple --extra-index-url https://download.pytorch.org/whl/cu128 sleap-nn track --data-path video.mp4 --model-paths model_ckpt_dir/

uvx --from "sleap-nn[torch]" --index-url https://pypi.org/simple --extra-index-url https://download.pytorch.org/whl/cpu sleap-nn track --data-path video.mp4 --model-paths model_ckpt_dir/

uvx sleap-nn[torch] track --data-path video.mp4 --model-paths model_ckpt_dir/

Info

  • For more information on which CUDA version to use for your system, see the PyTorch installation guide.
    The --extra-index-url in the install command should match the CUDA version you need (e.g., https://download.pytorch.org/whl/cuda118 for CUDA 11.8, https://download.pytorch.org/whl/cuda128 for CUDA 12.8, etc.).
  • On macOS, MPS (Metal Performance Shaders) is automatically enabled for Apple Silicon acceleration.

Model Paths

--model-paths should be set to <config.trainer_config.config_dir>/<config.trainer_config.run_name>. Make sure the model checkpoint directory contains both best.ckpt (or legacy sleap best_model.h5 - only UNet backbone is supported) and training_config.yaml (or legacy sleap training_config.json - only UNet backbone is supported) files. The inference will fail without these files.

📚 Documentation Structure

  • Installation Guide - Complete installation instructions with different options (CPU, GPU, development)

How-to Guides

  • Configuration Guide - Detailed explanation of all configuration parameters and how to set up your config file for training
  • Training Guide - How to train models using the CLI or Python API and advanced training options
  • Inference Guide - How to run inference and tracking with CLI/ APIs and evaluate the models

Tutorials

Additional Resources

  • Model Architectures — Explore supported model types, backbone networks, and architecture details.
  • Core Components — Main building blocks of the sleap-nn pipeline.
  • API Reference — Complete API documentation for all sleap-nn modules and functions.

Get Help

If you encounter issues or have questions about SLEAP-NN:

Report Bugs:

Found a bug? Please create an issue on GitHub: - Create a new issue - Include details about your environment, error messages, and steps to reproduce

Start a Discussion:

Have questions about usage, feature requests, or want to share your experience? - Start a discussion

Additional Resources:

Next Steps