SuperAnimal pretrained pose estimation models for behavioral analysis

TL;DR

This work presents SuperAnimal, a method to develop unified foundation models that can be used on over 45 species, without additional manual labels, that show excellent performance across six pose estimation benchmarks.

Abstract

Quantification of behavior is critical in applications ranging from neuroscience, veterinary medicine and animal conservation efforts. A common key step for behavioral analysis is first extracting relevant keypoints on animals, known as pose estimation. However, reliable inference of poses currently requires domain knowledge and manual labeling effort to build supervised models. We present a series of technical innovations that enable a new method, collectively called SuperAnimal, to develop unified foundation models that can be used on over 45 species, without additional human labels. Concretely, we introduce a method to unify the keypoint space across differently labeled datasets (via our generalized data converter) and for training these diverse datasets in a manner such that they don't catastrophically forget keypoints given the unbalanced inputs (via our keypoint gradient masking and memory replay approaches). These models show excellent performance across six pose benchmarks. Then, to ensure maximal usability for end-users, we demonstrate how to fine-tune the models on differently labeled data and provide tooling for unsupervised video adaptation to boost performance and decrease jitter across frames. If the models are fine-tuned, we show SuperAnimal models are 10-100$\times$ more data efficient than prior transfer-learning-based approaches. We illustrate the utility of our models in behavioral classification in mice and gait analysis in horses. Collectively, this presents a data-efficient solution for animal pose estimation.

Figures (11)

• Figure 1: The DeepLabCut Model Zoo, the SuperAnimal method, and SuperAnimal-TopViewMouse model performance.a: The website can collect data shared by the research community; SuperAnimal models are trained, and can be used for inference on novel images and videos with or without further training (fine-tuning). b: The panoptic animal pose estimation approach unifies the vocabulary of pose data across labs, such that each individual dataset is a subset of a super-set keypoint space, independently of its naming. c: For canonical task-agnostic transfer learning, the encoder learns universal visual features from ImageNet, and a randomly initialized decoder is used to learn the pose from the downstream dataset. For task-aware fine-tuning, both encoder and decoder learn task-related visual-pose features in the pre-training datasets and the decoder is fine-tuned to update pose priors in downstream datasets. Crucially, the network has pose-estimation specific weights. d: Memory replay combines the strengths of SuperAnimal models' zero-shot inference, data combination strategy, and leveraging labeled data for fine-tuning (if needed). e: Data efficiency of baseline (ImageNet) and various SuperAnimal fine-tuning methods using bottom-up DLCRNet on the DLC-Openfield OOD dataset. 1-100% of the train data is 10, 50, 101, 506, 1012 frames respectively. Blue shadow represents minimum, maximum and blue dash is the mean for zero-shot performance across three shuffles. Large, connected dots represent mean results across three shuffles and smaller dots represent results for individual shuffles. Inset: Using memory replay avoids catastrophic forgetting. f: SuperAnimal vs. baseline results on the TriMouse benchmark, showing zero-shot performance with top-down HRNet and AnimalTokenPose, and fine-tuning results with HRNet. 1-100% of the train data is 1, 7, 15, 76, 152 frames respectively Inset: example image of results. g: SuperAnimal-TopViewMouse (DLCRNet) qualitative results on the within distribution test images (IID). They were randomly selected based on the visibility of the keypoints within the figure (but not on performance). Full keypoint color and mapping is available in Extended Data Figure \ref{['fig:Extended_data_fig1']}). h: Visualization of model performance on OOD images using DLCRNet.
• Figure 2: SuperAnimal-Quadrupeda: Qualitative performance with SuperAnimal-Quadruped (HRNet-w32). Image randomly selected based on visibility of the keypoints within the figure (but not on performance). A likelihood cutoff of 0.6 was applied for keypoint visualization. Full keypoint color and mapping is available in Extended Data Figure \ref{['fig:Extended_data_fig1']}). HRNet-w32 and same cutoff of 0.6 are used in other panels. b: Performance on the official OOD Horse-10 test set, training with the official IID splits, reported as a normalized error from eye to nose, see inset adopted from mathis2021pretraining and qualitative zero-shot performance. HRNet-w32 is trained on AP-10K and Quadruped-80K, respectively, for zero-shot performance comparison. 1-100% of the data is 14, 73, 146, 734, 1469 frames, respectively. c: Performance on the OOD iRodent dataset, reported mAP. Colors and zero-shot baseline are as in b. 1-100% of the data is 3, 17, 35, 177, 354 frames, respectively. See inset for qualitative zero-shot performance. d: Performance on the OOD AnimalPose dataset, reported as mAP. HRNet-w32 trained on AP-10K is used as an additional zero-shot baseline. Qualitative zero-shot performance is shown. e: Performance on the OOD AP-10K dataset, reported as mAP. Qualitative zero-shot performance is also shown. f: AP-10K benchmark with SA-Q and other pose data pre-trained models. The size of dots represents the parameter size of each model. The number of pre-training images represents the number of pose data models trained before being fine-tuned on AP-10K.
• Figure 3: Unsupervised video adaptation methods.a: Illustration of the unsupervised video adaptation algorithm. b-e: Animal size described by convex hull of keypoints using the SA-TVM model. Frequent changes of the convex hull indicates non-smooth keypoint predictions, and below are example images with and without video adaptation showing the trailing keypoints for 10 past frames of data (to demonstrate the motion smoothness). f-g: Change in jitter score before and after video adaptation. Overall, our method had a significant effect on reducing jitter ($F(1,23286)$=190.03, $p$<.0001; Table \ref{['jitter_lme']}, in all but the dog ($p$=.36, $d$=-.03) and Golden lab ($p$=.62, $d$=-.06) videos; Table \ref{['contrasts_jitter']}. h-i: Same analysis as in b-e using the SA-Q model. Note that an additional median filtering post-video adaptation examples (dark purple line) can be used if needed. k: Video adaptation, self-pacing and Kalman filter's performance on the Horse-30 video dataset where j is an example of one of 30 videos from the dataset.
• Figure 4: Zero-shot behavioral quantification with SuperAnimal.a: Workflow overview for behavioral analysis with SuperAnimal. b: Images of the open-source dataset from Sturman et al. sturman2020deep with their DeepLabCut "in distribution" model and our SuperAnimal zero-shot, out-of-distribution, results. c: Ethogram comparing ground truth annotations vs. zero-shot predictions from SuperAnimal-TopViewMouse. d: F1 score computed across IID (Sturman) and SuperAnimal with, or without CEBRA on the two behavioral classes. e: CEBRA schneider2022cebra embedding on Sturman keypoints and SuperAnimal-based keypoints in 3D, transformed with FastICA for visualization. f: Correlation matrix that demonstrates the correlation between SuperAnimal-TopViewMouse and ground-truth annotations averaged across 3 annotators and across the model and keypoint configurations. g: We analyzed 30 horse videos where every frame had a ground truth (GT) annotation of keypoints mathis2021pretraining (left) vs. our SuperAnimal-Quadruped model (right). The right limbs (closest to the camera) from one representative gait trial are shown. Swing and stance phases are colored in light grey and black zones, respectively. h: Histogram delineating the number of videos where the ground contact by the hoof were identical to the GT vs. over or under counted by 1 stride (no error larger than 1 was found). i: We computed the error between the GT stride length vs. model prediction for the hoofs (i.e, right_back_paw vs. Offhindfoot, etc). Each dot represents a stride, color denotes hindlimb vs. forelimb, near legs only.
• Figure S1: Constructing SuperAnimal models and keypoint gradient masking.a: Demonstration of how multiple pose datasets are merged into a single dataset. We created a main keypoint names to cover all keypoints we observe from datasets. Then we built a conversion table to map keypoints from each dataset to the main keypoint names. We design a corresponding conversion table such that anatomically similar keypoints are mapped to the same keypoint. Below we add the keypoint naming map for both SuperAnimal-TopViewMouse and SuperAnimal-Quadruped models. b: Composition of the SuperAnimal-Quadruped (left) and SuperAnimal-TopViewMouse (right) datasets. c: Demonstration of keypoint gradient masking algorithm. Keypoints that were not defined in the original datasets introduce false penalties for the model training. Therefore, during back-propagation, the gradients of those undefined keypoints are artificially masked. d: With masking, the model is able to learn a pose representation that is the union of training datasets. Without masking, the model has severe degraded pose representation.