lightning_pose.models.heads

lightning_pose.models.heads.heatmap Module

Heads that produce heatmap predictions for heatmap regression.

Functions

make_upsampling_layers(in_channels, ...)
initialize_upsampling_layers(layers)	Intialize the Conv2DTranspose upsampling layers.
upsample(inputs)	Upsample batch of heatmaps by a factor of two using interpolation (no learned weights).
run_subpixelmaxima(heatmaps, ...)	Use soft argmax on heatmaps.

Classes

HeatmapHead

Simple deconvolution head that converts 2D feature maps to per-keypoint heatmaps.

lightning_pose.models.heads.heatmap_mhcrnn Module

Heads that produce heatmap predictions for heatmap regression.

Classes

HeatmapMHCRNNHead	Multi-head convolutional recurrent neural network head.
UpsamplingCRNN	Bidirectional Convolutional RNN network that handles heatmaps of context frames.

lightning_pose.models.heads.regression Module

Heads that produce (x, y) predictions for coordinate regression.

Classes

LinearRegressionHead

Linear regression head that converts 2D feature maps to a vector of (x, y) coordinates.