Spaces:
Sleeping
Sleeping
File size: 2,775 Bytes
b38c7b5 |
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 |
import torch
# batch*n
def normalize_vector(v):
batch = v.shape[0]
v_mag = torch.sqrt(v.pow(2).sum(1)) # batch
eps = torch.tensor(1e-8, device=v.device)
v_mag = torch.max(v_mag, eps)
v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
v = v / v_mag
return v
# u, v batch*n
def cross_product(u, v):
batch = u.shape[0]
# print (u.shape)
# print (v.shape)
i = u[:, 1] * v[:, 2] - u[:, 2] * v[:, 1]
j = u[:, 2] * v[:, 0] - u[:, 0] * v[:, 2]
k = u[:, 0] * v[:, 1] - u[:, 1] * v[:, 0]
out = torch.cat((i.view(batch, 1), j.view(batch, 1), k.view(batch, 1)), 1) # batch*3
return out
# poses batch*6
# poses
def compute_rotation_matrix_from_ortho6d(poses):
x_raw = poses[:, 0:3] # batch*3
y_raw = poses[:, 3:6] # batch*3
x = normalize_vector(x_raw) # batch*3
z = cross_product(x, y_raw) # batch*3
z = normalize_vector(z) # batch*3
y = cross_product(z, x) # batch*3
x = x.view(-1, 3, 1)
y = y.view(-1, 3, 1)
z = z.view(-1, 3, 1)
matrix = torch.cat((x, y, z), 2) # batch*3*3
return matrix
# input batch*4*4 or batch*3*3
# output torch batch*3 x, y, z in radiant
# the rotation is in the sequence of x,y,z
def compute_euler_angles_from_rotation_matrices(rotation_matrices):
batch = rotation_matrices.shape[0]
R = rotation_matrices
sy = torch.sqrt(R[:, 0, 0] * R[:, 0, 0] + R[:, 1, 0] * R[:, 1, 0])
singular = sy < 1e-6
singular = singular.float()
x = torch.atan2(R[:, 2, 1], R[:, 2, 2])
y = torch.atan2(-R[:, 2, 0], sy)
z = torch.atan2(R[:, 1, 0], R[:, 0, 0])
xs = torch.atan2(-R[:, 1, 2], R[:, 1, 1])
ys = torch.atan2(-R[:, 2, 0], sy)
zs = R[:, 1, 0] * 0
out_euler = torch.zeros(batch, 3, device=rotation_matrices.device)
out_euler[:, 0] = x * (1 - singular) + xs * singular
out_euler[:, 1] = y * (1 - singular) + ys * singular
out_euler[:, 2] = z * (1 - singular) + zs * singular
return out_euler
def get_R(x, y, z):
"""Get rotation matrix from three rotation angles (radians). right-handed.
Args:
x: rotation angle around x-axis
y: rotation angle around y-axis
z: rotation angle around z-axis
Returns:
R: [3, 3]. rotation matrix.
"""
# x
Rx = torch.tensor(
[[1, 0, 0], [0, torch.cos(x), -torch.sin(x)], [0, torch.sin(x), torch.cos(x)]],
device=x.device,
)
# y
Ry = torch.tensor(
[[torch.cos(y), 0, torch.sin(y)], [0, 1, 0], [-torch.sin(y), 0, torch.cos(y)]],
device=y.device,
)
# z
Rz = torch.tensor(
[[torch.cos(z), -torch.sin(z), 0], [torch.sin(z), torch.cos(z), 0], [0, 0, 1]],
device=z.device,
)
R = torch.mm(Rz, torch.mm(Ry, Rx))
return R
|