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

Parameters:

• data_module – give losses access to data for computing data-specific loss params

• epsilon – loss values below epsilon will be zeroed out

• log_weight – natural log of the weight in front of the loss term in the final objective function

__new__(**kwargs)