import numpy as np
from .mesh_eval import compute_similarity_transform
[docs]def keypoint_mpjpe(pred, gt, mask, alignment='none'):
"""Calculate the mean per-joint position error (MPJPE) and the error after
rigid alignment with the ground truth (P-MPJPE).
batch_size: N
num_keypoints: K
keypoint_dims: C
Args:
pred (np.ndarray[N, K, C]): Predicted keypoint location.
gt (np.ndarray[N, K, C]): Groundtruth keypoint location.
mask (np.ndarray[N, K]): Visibility of the target. False for invisible
joints, and True for visible. Invisible joints will be ignored for
accuracy calculation.
alignment (str, optional): method to align the prediction with the
groundtruth. Supported options are:
- ``'none'``: no alignment will be applied
- ``'scale'``: align in the least-square sense in scale
- ``'procrustes'``: align in the least-square sense in scale,
rotation and translation.
Returns:
tuple: A tuple containing joint position errors
- mpjpe (float|np.ndarray[N]): mean per-joint position error.
- p-mpjpe (float|np.ndarray[N]): mpjpe after rigid alignment with the
ground truth
"""
assert mask.any()
if alignment == 'none':
pass
elif alignment == 'procrustes':
pred = np.stack([
compute_similarity_transform(pred_i, gt_i)
for pred_i, gt_i in zip(pred, gt)
])
elif alignment == 'scale':
pred_dot_pred = np.einsum('nkc,nkc->n', pred, pred)
pred_dot_gt = np.einsum('nkc,nkc->n', pred, gt)
scale_factor = pred_dot_gt / pred_dot_pred
pred = pred * scale_factor[:, None, None]
else:
raise ValueError(f'Invalid value for alignment: {alignment}')
error = np.linalg.norm(pred - gt, ord=2, axis=-1)[mask].mean()
return error