Modified function 'apply_umeyama' in mvn/utils/tred.py

[RichardChen20]

I have modified the function 'apply_umeyama' so that it can directly process data from a batch much more efficiently. Here are the codes:

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)
gt_centered = batch_gt - batch_gt.mean(axis=1, keepdim=True)

H = pred_centered.transpose(1, 2).matmul(gt_centered)
u, _, v = torch.svd(H)  # Kabsch algorithm
R = v.matmul(u.transpose(1, 2))

if scaling:
    c = (gt_centered.norm(p='fro', dim=2) / pred_centered.norm(p='fro', dim=2)).mean(axis=1, keepdim=True).unsqueeze(-1)
else:
    c = 1.0

if rotation:
    return batch_pred.matmul(R.transpose(1, 2)) * c
else:
    return batch_pred * c