Spaces:
Runtime error
Runtime error
File size: 4,633 Bytes
da48dbe 487ee6d da48dbe fb140f6 da48dbe fb140f6 da48dbe fb140f6 da48dbe fb140f6 da48dbe fb140f6 da48dbe fb140f6 da48dbe fb140f6 da48dbe fb140f6 da48dbe fb140f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
"""
Parts of the code are adapted from https://github.com/akanazawa/hmr
"""
from __future__ import absolute_import, division, print_function
import numpy as np
import torch
def compute_similarity_transform(S1, S2):
"""
Computes a similarity transform (sR, t) that takes
a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
i.e. solves the orthogonal Procrutes problem.
"""
transposed = False
if S1.shape[0] != 3 and S1.shape[0] != 2:
S1 = S1.T
S2 = S2.T
transposed = True
assert (S2.shape[1] == S1.shape[1])
# 1. Remove mean.
mu1 = S1.mean(axis=1, keepdims=True)
mu2 = S2.mean(axis=1, keepdims=True)
X1 = S1 - mu1
X2 = S2 - mu2
# 2. Compute variance of X1 used for scale.
var1 = np.sum(X1**2)
# 3. The outer product of X1 and X2.
K = X1.dot(X2.T)
# 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
# singular vectors of K.
U, s, Vh = np.linalg.svd(K)
V = Vh.T
# Construct Z that fixes the orientation of R to get det(R)=1.
Z = np.eye(U.shape[0])
Z[-1, -1] *= np.sign(np.linalg.det(U.dot(V.T)))
# Construct R.
R = V.dot(Z.dot(U.T))
# 5. Recover scale.
scale = np.trace(R.dot(K)) / var1
# 6. Recover translation.
t = mu2 - scale * (R.dot(mu1))
# 7. Error:
S1_hat = scale * R.dot(S1) + t
if transposed:
S1_hat = S1_hat.T
return S1_hat
def compute_similarity_transform_batch(S1, S2):
"""Batched version of compute_similarity_transform."""
S1_hat = np.zeros_like(S1)
for i in range(S1.shape[0]):
S1_hat[i] = compute_similarity_transform(S1[i], S2[i])
return S1_hat
def reconstruction_error(S1, S2, reduction='mean'):
"""Do Procrustes alignment and compute reconstruction error."""
S1_hat = compute_similarity_transform_batch(S1, S2)
re = np.sqrt(((S1_hat - S2)**2).sum(axis=-1)).mean(axis=-1)
if reduction == 'mean':
re = re.mean()
elif reduction == 'sum':
re = re.sum()
return re, S1_hat
# https://math.stackexchange.com/questions/382760/composition-of-two-axis-angle-rotations
def axis_angle_add(theta, roll_axis, alpha):
"""Composition of two axis-angle rotations (PyTorch version)
Args:
theta: N x 3
roll_axis: N x 3
alph: N x 1
Returns:
equivalent axis-angle of the composition
"""
alpha = alpha / 2.
l2norm = torch.norm(theta + 1e-8, p=2, dim=1)
angle = torch.unsqueeze(l2norm, -1)
normalized = torch.div(theta, angle)
angle = angle * 0.5
b_cos = torch.cos(angle).cpu()
b_sin = torch.sin(angle).cpu()
a_cos = torch.cos(alpha)
a_sin = torch.sin(alpha)
dot_mm = torch.sum(normalized * roll_axis, dim=1, keepdim=True)
cross_mm = torch.zeros_like(normalized)
cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1]
cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2]
cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0]
c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm
c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm
c_angle = 2 * torch.acos(c_cos)
c_sin = torch.sin(c_angle * 0.5)
c_n = (c_angle / c_sin) * c_sin_n
return c_n
def axis_angle_add_np(theta, roll_axis, alpha):
"""Composition of two axis-angle rotations (NumPy version)
Args:
theta: N x 3
roll_axis: N x 3
alph: N x 1
Returns:
equivalent axis-angle of the composition
"""
alpha = alpha / 2.
angle = np.linalg.norm(theta + 1e-8, ord=2, axis=1, keepdims=True)
normalized = np.divide(theta, angle)
angle = angle * 0.5
b_cos = np.cos(angle)
b_sin = np.sin(angle)
a_cos = np.cos(alpha)
a_sin = np.sin(alpha)
dot_mm = np.sum(normalized * roll_axis, axis=1, keepdims=True)
cross_mm = np.zeros_like(normalized)
cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1]
cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2]
cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0]
c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm
c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm
c_angle = 2 * np.arccos(c_cos)
c_sin = np.sin(c_angle * 0.5)
c_n = (c_angle / c_sin) * c_sin_n
return c_n
|