generate_heatmaps

lightning_pose.data.utils.generate_heatmaps(keypoints: Tensor[Tensor], height: int, width: int, output_shape: Tuple[int, int], sigma: float | int = 1.25, uniform_heatmaps: bool = False) → Tensor[Tensor]

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

Returns:

batch of 2D heatmaps