generate_heatmaps

lightning_pose.data.utils.generate_heatmaps(keypoints: Float[Tensor, 'batch num_keypoints 2'], height: int, width: int, output_shape: tuple[int, int], sigma: float = 1.25, uniform_heatmaps: bool = False, keep_gradients: bool = False) → Float[Tensor, 'batch num_keypoints height width']

Generate 2D Gaussian heatmaps from mean and sigma.

Parameters:

• keypoints – coordinates that serve as mean of gaussian bump

• height – height of reshaped image (pixels, e.g., 128, 256, 512…)

• width – width of reshaped image (pixels, e.g., 128, 256, 512…)

• output_shape – dimensions of downsampled heatmap, (height, width)

• sigma – control spread of gaussian

• uniform_heatmaps – output uniform heatmaps if missing ground truth label, rather than skip

• keep_gradients – True to not detach gradients from keypoints before creating heatmaps

Returns:

batch of 2D heatmaps