自己写了一个3D human pose里面计算NMPJPE的函数,基于Pytorch,输入的pred和gt维度都是(batchsize, num_joints, 3)。
所谓NMPJPE就是先把pred经过旋转、放缩变换到和gt相同的尺度下,然后再计算两者之间关节点平均欧式距离,这个在unsupervised 3d human pose estimation中常用到。代码在下文,欢迎大家随意取用。
import torch
def apply_umeyama(batch_gt, batch_pred, rotation=True, scaling=True):
# pred and gt size (batch_size, num_joints, 3)
pred_centered = batch_pred - batch_pred.mean(axis=1, keepdim=True