Training and inference on an example dataset

In this notebook we’ll install SLEAP, download a sample dataset, run training and inference on that dataset using the SLEAP command-line interface, and then download the predictions.

Install SLEAP

Note: Before installing SLEAP check SLEAP releases page for the latest version.

!pip uninstall -y opencv-python opencv-contrib-python
!pip install sleap

Download sample training data into Colab

Let’s download a sample dataset from the SLEAP sample datasets repository into Colab.

!apt-get install tree
!wget -O dataset.zip https://github.com/talmolab/sleap-datasets/releases/download/dm-courtship-v1/drosophila-melanogaster-courtship.zip
!mkdir dataset
!unzip dataset.zip -d dataset
!rm dataset.zip
!tree dataset

Train models

For the top-down pipeline, we’ll need train two models: a centroid model and a centered-instance model.

Using the command-line interface, we’ll first train a model for centroids using the default training profile. The training profile determines the model architecture, the learning rate, and other parameters.

When you start training, you’ll first see the training parameters and then the training and validation loss for each training epoch.

As soon as you’re satisfied with the validation loss you see for an epoch during training, you’re welcome to stop training by clicking the stop button. The version of the model with the lowest validation loss is saved during training, and that’s what will be used for inference.

If you don’t stop training, it will run for 200 epochs or until validation loss fails to improve for some number of epochs (controlled by the early_stopping fields in the training profile).

!sleap-train baseline.centroid.json "dataset/drosophila-melanogaster-courtship/courtship_labels.slp" --run_name "courtship.centroid" --video-paths "dataset/drosophila-melanogaster-courtship/20190128_113421.mp4"

Let’s now train a centered-instance model.

!sleap-train baseline_medium_rf.topdown.json "dataset/drosophila-melanogaster-courtship/courtship_labels.slp" --run_name "courtship.topdown_confmaps" --video-paths "dataset/drosophila-melanogaster-courtship/20190128_113421.mp4"

The models (along with the profiles and ground truth data used to train and validate the model) are saved in the models/ directory:

!tree models/

models/
├── courtship.centroid
│   ├── best_model.h5
│   ├── initial_config.json
│   ├── labels_gt.train.slp
│   ├── labels_gt.val.slp
│   ├── training_config.json
│   └── training_log.csv
└── courtship.topdown_confmaps
    ├── best_model.h5
    ├── initial_config.json
    ├── labels_gt.train.slp
    ├── labels_gt.val.slp
    ├── training_config.json
    └── training_log.csv

2 directories, 12 files

Inference

Let’s run inference with our trained models for centroids and centered instances.

!sleap-track "dataset/drosophila-melanogaster-courtship/20190128_113421.mp4" --frames 0-100 -m "models/courtship.centroid" -m "models/courtship.topdown_confmaps"

When inference is finished, predictions are saved in a file. Since we didn’t specify a path, it will be saved as <video filename>.predictions.slp in the same directory as the video:

!tree dataset/drosophila-melanogaster-courtship

dataset/drosophila-melanogaster-courtship
├── 20190128_113421.mp4
├── 20190128_113421.mp4.predictions.slp
├── courtship_labels.slp
└── example.jpg

0 directories, 4 files

You can inspect your predictions file using sleap-inspect:

!sleap-inspect dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp

If you’re using Chrome you can download your trained models like so:

# Zip up the models directory
!zip -r trained_models.zip models/

# Download.
from google.colab import files
files.download("/content/trained_models.zip")

And you can likewise download your predictions:

from google.colab import files
files.download('dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp')

In some other browsers (Safari) you might get an error and you can instead download using the “Files” tab in the side panel (it has a folder icon). Select “Show table of contents” in the “View” menu if you don’t see the side panel.