ReprojectionHeatmapLoss

class lightning_pose.losses.losses.ReprojectionHeatmapLoss

Bases: Loss

Penalize error between predicted 2D->3D->2D->heatmap and ground truth heatmap.

Attributes Summary

loss_name

Methods Summary

__call__(heatmaps_targ, ...[, stage])	Call self as a function.
compute_loss(targets, predictions)
remove_nans(loss, targets)

Attributes Documentation

loss_name = 'supervised_reprojection_heatmap_mse'

Methods Documentation

__call__(heatmaps_targ: Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 'heatmap_height', 'heatmap_width'), 'cls_name': 'TensorType'}], keypoints_pred_2d_reprojected: Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 2), 'cls_name': 'TensorType'}], stage: Literal['train', 'val', 'test'] | None = None, **kwargs) → Tensor, {'__torchtyping__': True, 'details': ((),), 'cls_name': 'TensorType'}], list[dict]]

Call self as a function.

compute_loss(targets: Tensor, {'__torchtyping__': True, 'details': ('batch_x_num_keypoints', 'heatmap_height', 'heatmap_width'), 'cls_name': 'TensorType'}], predictions: Tensor, {'__torchtyping__': True, 'details': ('batch_x_num_keypoints', 'heatmap_height', 'heatmap_width'), 'cls_name': 'TensorType'}]) → Tensor, {'__torchtyping__': True, 'details': ('batch_x_num_keypoints', 'heatmap_height', 'heatmap_width',), 'cls_name': 'TensorType'}]

remove_nans(loss: Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 'heatmap_height', 'heatmap_width'), 'cls_name': 'TensorType'}], targets: Tensor, {'__torchtyping__': True, 'details': ('batch', 'num_keypoints', 'heatmap_height', 'heatmap_width'), 'cls_name': 'TensorType'}]) → Tensor, {'__torchtyping__': True, 'details': ('valid_losses',), 'cls_name': 'TensorType'}]

__init__(original_image_height: int, original_image_width: int, downsampled_image_height: int, downsampled_image_width: int, log_weight: float = 0.0, uniform_heatmaps: bool = False, **kwargs) → None

Initialize ReprojectionHeatmapLoss.

Converts 2D reprojected keypoints (obtained by projecting 3D triangulated predictions back into each camera’s image plane) into heatmaps and compares them with the ground truth heatmaps using pixel-wise MSE.

Parameters:

• original_image_height – height of the full-resolution input image in pixels.

• original_image_width – width of the full-resolution input image in pixels.

• downsampled_image_height – height of the heatmap output (after backbone downsampling).

• downsampled_image_width – width of the heatmap output.

• log_weight – final weight in front of the loss term in the objective function is computed as 1.0 / (2.0 * exp(log_weight)).

• uniform_heatmaps – if True, generate uniform (flat) target heatmaps for NaN ground truth keypoints instead of ignoring them in the loss.

__new__(**kwargs)