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)

Methods Documentation

forward(features: Tensor, {'__torchtyping__': True, 'details': ('batch', 'features', 'features_height', 'features_width'), 'cls_name': 'TensorType'}]) → Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 'heatmap_height', 'heatmap_width',), 'cls_name': 'TensorType'}]

Upsample representations and normalize to get final heatmaps.

run_subpixelmaxima(heatmaps)

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

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)