HeatmapHead

class lightning_pose.models.heads.HeatmapHead

Bases: Module

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

This is the standard heatmap head used in the Lightning Pose package. The head is composed of several deconvolutional layers followed by a 2D spatial softmax to generate normalized heatmaps from low-resolution feature maps.

Methods Summary

forward(features)	Upsample representations and normalize to get final heatmaps.
run_subpixelmaxima(heatmaps)	Apply soft argmax to heatmaps to obtain subpixel keypoint predictions.

Methods Documentation

forward(features: Float[Tensor, 'batch features features_height features_width']) → Float[Tensor, 'batch num_keypoints heatmap_height heatmap_width']

Upsample representations and normalize to get final heatmaps.

run_subpixelmaxima(heatmaps: Float[Tensor, 'batch num_keypoints height width']) → tuple[Tensor, Tensor]

Apply soft argmax to heatmaps to obtain subpixel keypoint predictions.

Parameters:

heatmaps – predicted heatmaps of shape (batch, num_keypoints, height, width).

Returns:

Tuple of (keypoints, confidences) where keypoints has shape (batch, num_targets) and confidences has shape (batch, num_keypoints).

__init__(backbone_arch: str, in_channels: int, out_channels: int, deconv_out_channels: int | None = None, downsample_factor: int = 2, final_softmax: bool = True) → None

Parameters:

• backbone_arch – string denoting backbone architecture; to remove in future release

• in_channels – number of channels in the input feature map

• out_channels – number of channels in the output heatmap (i.e. number of keypoints)

• deconv_out_channels – output channel number for each intermediate deconv layer; defaults to number of keypoints

• downsample_factor – make heatmaps smaller than input frames by this factor; subpixel operations are performed for increased precision

• final_softmax – pass final heatmaps through a 2D softmax with temperature 1.0

__new__(**kwargs)