evaluate_heatmaps_at_location

lightning_pose.data.utils.evaluate_heatmaps_at_location(heatmaps: Float[Tensor, 'batch num_keypoints heatmap_height heatmap_width'], locs: Float[Tensor, 'batch num_keypoints 2'], sigma: float = 1.25, num_stds: int = 2) → Float[Tensor, 'batch num_keypoints']

Evaluate 4D heatmaps using a 3D location tensor (last dim is x, y coords). Since the model outputs heatmaps with a standard deviation of sigma, confidence will be spread across neighboring pixels. To account for this, confidence is computed by taking all pixels within two standard deviations of the predicted pixel.