HeatmapMHCRNNHead

class lightning_pose.models.heads.HeatmapMHCRNNHead

Bases: Module

Multi-head convolutional recurrent neural network head.

This head converts a sequence of 2D feature maps to per-keypoint heatmaps for the center frame. The head is composed of two heads: - single frame head: several deconvolutional layers followed by a 2D spatial softmax to

generate normalized heatmaps from low-resolution feature maps for a single frame.

  • multi-frame head: several deconvolutional layers are applied to each set of features in a temporal sequence; the resulting heatmaps are fed into a convolutional recurrent neural network to produce heatmaps for the center frame

Methods Summary

forward(features,Β batch_shape,Β is_multiview)

Handle context frames then upsample 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 rep_height rep_width frames'], batch_shape: Tensor, is_multiview: bool) → tuple[Float[Tensor, 'batch num_keypoints heatmap_height heatmap_width'], Float[Tensor, 'batch num_keypoints heatmap_height heatmap_width']][source]

Handle context frames then upsample to get final heatmaps.

Parameters:

  • features – outputs of backbone

  • batch_shape – identifies whether or not we need to do some reshaping

  • is_multiview – if batch has a view dimension

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

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, upsampling_factor: Literal[1, 2] = 2) → None[source]
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

  • upsampling_factor – upsample features before feeding to crnn

__new__(**kwargs)