Spaces:
Sleeping
Sleeping
| 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 | |