Neural network backend for training and inference for animal pose estimation.
This 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.
📚 Documentation - Comprehensive guides and API reference
Prerequisites: Python 3.11
- Windows/Linux with NVIDIA GPU (CUDA 11.8):
pip install sleap-nn[torch-cuda118]- Windows/Linux with NVIDIA GPU (CUDA 12.8):
pip install sleap-nn[torch-cuda128]- macOS with Apple Silicon (M1, M2, M3, M4) or CPU-only (no GPU or unsupported GPU): Note: Even if torch-cpu is used on macOS, the MPS backend will be available.
pip install sleap-nn[torch-cpu]- Clone the sleap-nn repo
git clone https://github.com/talmolab/sleap-nn.git
cd sleap-nn-
Install
uvand development dependencies
uvis a fast and modern package manager forpyproject.toml-based projects. Refer installation docs to install uv. -
Install sleap-nn dependencies based on your platform\
-
Sync all dependencies based on your correct wheel using
uv sync:- Windows/Linux with NVIDIA GPU (CUDA 11.8):
uv sync --extra dev --extra torch-cuda118
- Windows/Linux with NVIDIA GPU (CUDA 12.8):
uv sync --extra dev --extra torch-cuda128
- macOS with Apple Silicon (M1, M2, M3, M4) or CPU-only (no GPU or unsupported GPU): Note: Even if torch-cpu is used on macOS, the MPS backend will be available.
uv sync --extra dev --extra torch-cpu
You can find the correct wheel for your system at:
👉 https://pytorch.org/get-started/locally -
-
Run tests
uv run pytest tests
-
(Optional) Lint and format code
uv run black --check sleap_nn tests uv run ruff check sleap_nn/
Let's start SLEAPiNNg !!! 🐭🐭
For detailed information on setting up config, training/ inference workflows, please refer to our docs.
Create a config.yaml file for your experiment.
Use a sample config from
docs/sample_configs.
Download sample training data from here and validation data from here for quick experimentation.
sleap-nn-train --config-name config.yaml --config-dir configs/ "data_config.train_labels_path=[labels.pkg.slp]"To run inference:
sleap-nn-track --data-path video.mp4 --model-paths model_ckpt_dir/