evaluation_matrix.py

import torch
import numpy as np

def mpjpe(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape)-1))

def weighted_mpjpe(predicted, target, w):
    """
    Weighted mean per-joint position error (i.e. mean Euclidean distance)
    """
    assert predicted.shape == target.shape
    assert w.shape[0] == predicted.shape[0]
    return torch.mean(w * torch.norm(predicted - target, dim=len(target.shape)-1))

def p_mpjpe_torch(predicted, target, with_sRt=False,full_torch=False,with_aligned=False):
    """
    Pose error: MPJPE after rigid alignment (scale, rotation, and translation),
    often referred to as "Protocol #2" in many papers.
    """
    assert predicted.shape == target.shape

    muX = torch.mean(target, dim=1, keepdim=True)
    muY = torch.mean(predicted, dim=1, keepdim=True)

    X0 = target - muX
    Y0 = predicted - muY
    X0[X0**2<1e-6]=1e-3

    normX = torch.sqrt(torch.sum(X0**2, dim=(1, 2), keepdim=True))
    normY = torch.sqrt(torch.sum(Y0**2, dim=(1, 2), keepdim=True))

    normX[normX<1e-3]=1e-3

    X0 /= normX
    Y0 /= normY

    H = torch.matmul(X0.transpose(1,2), Y0)
    if full_torch:
        U, s, V = batch_svd(H)
    else:
        U, s, Vt = np.linalg.svd(H.cpu().numpy())
        V = torch.from_numpy(Vt.transpose(0, 2, 1)).cuda()
        U = torch.from_numpy(U).cuda()
        s = torch.from_numpy(s).cuda()

    R = torch.matmul(V, U.transpose(2, 1))

    # Avoid improper rotations (reflections), i.e. rotations with det(R) = -1
    sign_detR = torch.sign(torch.unsqueeze(torch.det(R[0]), 0))
    V[:, :, -1] *= sign_detR.unsqueeze(0)
    s[:, -1] *= sign_detR.flatten()
    R = torch.matmul(V, U.transpose(2, 1)) # Rotation

    tr = torch.unsqueeze(torch.sum(s, dim=1, keepdim=True), 2)

    a = tr * normX / normY # Scale
    t = muX - a*torch.matmul(muY, R) # Translation

    if (a!=a).sum()>0:
        print('NaN Error!!')
        print('UsV:',U,s,V)
        print('aRt:',a,R,t)
    a[a!=a]=1.
    R[R!=R]=0.
    t[t!=t]=0.
    # Perform rigid transformation on the input
    predicted_aligned = a*torch.matmul(predicted, R) + t
    if with_sRt:
        return torch.sqrt(((predicted_aligned - target)**2).sum(-1)).mean(),(a,R,t)#torch.mean(torch.norm(predicted_aligned - target, dim=len(target.shape)-1))
    if with_aligned:
        return torch.sqrt(((predicted_aligned - target)**2).sum(-1)).mean(),predicted_aligned
    # Return MPJPE
    return torch.sqrt(((predicted_aligned - target)**2).sum(-1)).mean()#torch.mean(torch.norm(predicted_aligned - target, dim=len(target.shape)-1))#,(a,R,t),predicted_aligned

def batch_svd(H):
    num = H.shape[0]
    U_batch, s_batch, V_batch = [],[],[]
    for i in range(num):
        U, s, V = H[i].svd(some=False)
        U_batch.append(U.unsqueeze(0))
        s_batch.append(s.unsqueeze(0))
        V_batch.append(V.unsqueeze(0))
    return torch.cat(U_batch,0),torch.cat(s_batch,0),torch.cat(V_batch,0)

def p_mpjpe(predicted, target, with_sRt=False,full_torch=False,with_aligned=False,each_separate=False):
    """
    Pose error: MPJPE after rigid alignment (scale, rotation, and translation),
    often referred to as "Protocol #2" in many papers.
    """
    assert predicted.shape == target.shape

    muX = np.mean(target, axis=1, keepdims=True)
    muY = np.mean(predicted, axis=1, keepdims=True)

    X0 = target - muX
    Y0 = predicted - muY

    normX = np.sqrt(np.sum(X0**2, axis=(1, 2), keepdims=True))
    normY = np.sqrt(np.sum(Y0**2, axis=(1, 2), keepdims=True))

    X0 /= (normX+1e-6)
    Y0 /= (normY+1e-6)


    H = np.matmul(X0.transpose(0, 2, 1), Y0).astype(np.float16).astype(np.float64)
    U, s, Vt = np.linalg.svd(H)
    V = Vt.transpose(0, 2, 1)
    R = np.matmul(V, U.transpose(0, 2, 1))

    # Avoid improper rotations (reflections), i.e. rotations with det(R) = -1
    sign_detR = np.sign(np.expand_dims(np.linalg.det(R), axis=1))
    V[:, :, -1] *= sign_detR
    s[:, -1] *= sign_detR.flatten()
    R = np.matmul(V, U.transpose(0, 2, 1)) # Rotation

    tr = np.expand_dims(np.sum(s, axis=1, keepdims=True), axis=2)

    a = tr * normX / normY # Scale
    t = muX - a*np.matmul(muY, R) # Translation

    # Perform rigid transformation on the input
    predicted_aligned = a*np.matmul(predicted, R) + t
    if each_separate:
        return np.linalg.norm(predicted_aligned - target, axis=len(target.shape)-1)

    error = np.mean(np.linalg.norm(predicted_aligned - target, axis=len(target.shape)-1))
    if with_sRt and not with_aligned:
        return error, (a,R,t)
    if with_aligned:
        return error,(a,R,t),predicted_aligned
    # Return MPJPE
    return error

def n_mpjpe(predicted, target):
    """
    Normalized MPJPE (scale only), adapted from:
    https://github.com/hrhodin/UnsupervisedGeometryAwareRepresentationLearning/blob/master/losses/poses.py
    """
    assert predicted.shape == target.shape

    norm_predicted = torch.mean(torch.sum(predicted**2, dim=3, keepdim=True), dim=2, keepdim=True)
    norm_target = torch.mean(torch.sum(target*predicted, dim=3, keepdim=True), dim=2, keepdim=True)
    scale = norm_target / norm_predicted
    return mpjpe(scale * predicted, target)

def mean_velocity_error(predicted, target):
    """
    Mean per-joint velocity error (i.e. mean Euclidean distance of the 1st derivative)
    """
    assert predicted.shape == target.shape

    velocity_predicted = np.diff(predicted, axis=0)
    velocity_target = np.diff(target, axis=0)

    return np.mean(np.linalg.norm(velocity_predicted - velocity_target, axis=len(target.shape)-1))

def test():
    r1 = np.random.rand(3,14,3)
    r2 = np.random.rand(3,14,3)
    pmpjpe = p_mpjpe(r1, r2,with_sRt=False)
    pmpjpe_torch = p_mpjpe_torch(torch.from_numpy(r1), torch.from_numpy(r2),with_sRt=False,full_torch=True)
    print('pmpjpe: {}; {:.6f}; {:.6f}; {:.6f}'.format(np.abs(pmpjpe-pmpjpe_torch.numpy())<0.01,pmpjpe,pmpjpe_torch.numpy(), pmpjpe-pmpjpe_torch.numpy()))

if __name__ == '__main__':
    test()