Pose estimation made better, easier, and faster with video semi-supervised learning on the cloud

Computer vision algorithms now allow neuroscientists to quantify animal movement in unprecedented detail, shedding new light on problems in motor control, social behavior, navigation, and beyond. Existing models are only trained with labeled frames (supervised learning). Although effective in many cases, the supervised approach requires extensive image annotation, struggles to generalize to new videos, and produces noisy outputs that hinder downstream analyses. We address each of these limitations by introducing unlabeled video clips to model training and inference, and exploiting their spatiotemporal statistics in two different ways. First, we develop a suite of unsupervised training objectives that penalize the network whenever its predictions violate multiple-view geometry, smoothness of physical motion, or depart from a low-dimensional subspace of plausible body configurations. Second, we develop a new network architecture that predicts pose for a given frame using temporal context from surrounding (unlabeled) frames. Context frames resolve most brief occlusions and detection ambiguities (such as confusion between nearby and similar-looking body parts). We present results on two datasets: freely swimming mormyrid fish and head-fixed running mice, both densely tracked from multiple views. Our algorithms achieve better performance with fewer labels (sample-efficient), generalize better to unseen videos even when labels are abundant (out-of-distribution robustness), and provide smoother and more reliable trajectories for downstream analysis. Finally, we release a PyTorch Lightning package that supports easy model development and accelerated training, and a cloud-hosted application that allows users to annotate data, quickly train models, and predict new videos, directly from the browser and without any local installation.