generate_heatmaps

lightning_pose.data.utils.generate_heatmaps(keypoints: Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 2), 'cls_name': 'TensorType'}], height: int, width: int, output_shape: tuple[int, int], sigma: float = 1.25, uniform_heatmaps: bool = False, keep_gradients: bool = False) → Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 'height', 'width',), 'cls_name': 'TensorType'}][source]

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