evaluate_heatmaps_at_location

lightning_pose.data.utils.evaluate_heatmaps_at_location(heatmaps: Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 'heatmap_height', 'heatmap_width'), 'cls_name': 'TensorType'}], locs: Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 2), 'cls_name': 'TensorType'}], sigma: float = 1.25, num_stds: int = 2) → Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints',), 'cls_name': 'TensorType'}]

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.