Spaces:
Sleeping
Sleeping
| # Copyright (c) 2018-present, Facebook, Inc. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| # | |
| import torch | |
| import numpy as np | |
| _EPS4 = np.finfo(float).eps * 4.0 | |
| _FLOAT_EPS = np.finfo(np.float).eps | |
| # PyTorch-backed implementations | |
| def qinv(q): | |
| assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)' | |
| mask = torch.ones_like(q) | |
| mask[..., 1:] = -mask[..., 1:] | |
| return q * mask | |
| def qinv_np(q): | |
| assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)' | |
| return qinv(torch.from_numpy(q).float()).numpy() | |
| def qnormalize(q): | |
| assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)' | |
| return q / torch.norm(q, dim=-1, keepdim=True) | |
| def qmul(q, r): | |
| """ | |
| Multiply quaternion(s) q with quaternion(s) r. | |
| Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions. | |
| Returns q*r as a tensor of shape (*, 4). | |
| """ | |
| assert q.shape[-1] == 4 | |
| assert r.shape[-1] == 4 | |
| original_shape = q.shape | |
| # Compute outer product | |
| terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4)) | |
| w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3] | |
| x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2] | |
| y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1] | |
| z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0] | |
| return torch.stack((w, x, y, z), dim=1).view(original_shape) | |
| def qrot(q, v): | |
| """ | |
| Rotate vector(s) v about the rotation described by quaternion(s) q. | |
| Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v, | |
| where * denotes any number of dimensions. | |
| Returns a tensor of shape (*, 3). | |
| """ | |
| assert q.shape[-1] == 4 | |
| assert v.shape[-1] == 3 | |
| assert q.shape[:-1] == v.shape[:-1] | |
| original_shape = list(v.shape) | |
| # print(q.shape) | |
| q = q.contiguous().view(-1, 4) | |
| v = v.contiguous().view(-1, 3) | |
| qvec = q[:, 1:] | |
| uv = torch.cross(qvec, v, dim=1) | |
| uuv = torch.cross(qvec, uv, dim=1) | |
| return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape) | |
| def qeuler(q, order, epsilon=0, deg=True): | |
| """ | |
| Convert quaternion(s) q to Euler angles. | |
| Expects a tensor of shape (*, 4), where * denotes any number of dimensions. | |
| Returns a tensor of shape (*, 3). | |
| """ | |
| assert q.shape[-1] == 4 | |
| original_shape = list(q.shape) | |
| original_shape[-1] = 3 | |
| q = q.view(-1, 4) | |
| q0 = q[:, 0] | |
| q1 = q[:, 1] | |
| q2 = q[:, 2] | |
| q3 = q[:, 3] | |
| if order == 'xyz': | |
| x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) | |
| y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1 + epsilon, 1 - epsilon)) | |
| z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3)) | |
| elif order == 'yzx': | |
| x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3)) | |
| y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3)) | |
| z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1 + epsilon, 1 - epsilon)) | |
| elif order == 'zxy': | |
| x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1 + epsilon, 1 - epsilon)) | |
| y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) | |
| z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q1 * q1 + q3 * q3)) | |
| elif order == 'xzy': | |
| x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3)) | |
| y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3)) | |
| z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1 + epsilon, 1 - epsilon)) | |
| elif order == 'yxz': | |
| x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1 + epsilon, 1 - epsilon)) | |
| y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2 * (q1 * q1 + q2 * q2)) | |
| z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 * q1 + q3 * q3)) | |
| elif order == 'zyx': | |
| x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) | |
| y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1 + epsilon, 1 - epsilon)) | |
| z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3)) | |
| else: | |
| raise | |
| if deg: | |
| return torch.stack((x, y, z), dim=1).view(original_shape) * 180 / np.pi | |
| else: | |
| return torch.stack((x, y, z), dim=1).view(original_shape) | |
| # Numpy-backed implementations | |
| def qmul_np(q, r): | |
| q = torch.from_numpy(q).contiguous().float() | |
| r = torch.from_numpy(r).contiguous().float() | |
| return qmul(q, r).numpy() | |
| def qrot_np(q, v): | |
| q = torch.from_numpy(q).contiguous().float() | |
| v = torch.from_numpy(v).contiguous().float() | |
| return qrot(q, v).numpy() | |
| def qeuler_np(q, order, epsilon=0, use_gpu=False): | |
| if use_gpu: | |
| q = torch.from_numpy(q).cuda().float() | |
| return qeuler(q, order, epsilon).cpu().numpy() | |
| else: | |
| q = torch.from_numpy(q).contiguous().float() | |
| return qeuler(q, order, epsilon).numpy() | |
| def qfix(q): | |
| """ | |
| Enforce quaternion continuity across the time dimension by selecting | |
| the representation (q or -q) with minimal distance (or, equivalently, maximal dot product) | |
| between two consecutive frames. | |
| Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints. | |
| Returns a tensor of the same shape. | |
| """ | |
| assert len(q.shape) == 3 | |
| assert q.shape[-1] == 4 | |
| result = q.copy() | |
| dot_products = np.sum(q[1:] * q[:-1], axis=2) | |
| mask = dot_products < 0 | |
| mask = (np.cumsum(mask, axis=0) % 2).astype(bool) | |
| result[1:][mask] *= -1 | |
| return result | |
| def euler2quat(e, order, deg=True): | |
| """ | |
| Convert Euler angles to quaternions. | |
| """ | |
| assert e.shape[-1] == 3 | |
| original_shape = list(e.shape) | |
| original_shape[-1] = 4 | |
| e = e.view(-1, 3) | |
| ## if euler angles in degrees | |
| if deg: | |
| e = e * np.pi / 180. | |
| x = e[:, 0] | |
| y = e[:, 1] | |
| z = e[:, 2] | |
| rx = torch.stack((torch.cos(x / 2), torch.sin(x / 2), torch.zeros_like(x), torch.zeros_like(x)), dim=1) | |
| ry = torch.stack((torch.cos(y / 2), torch.zeros_like(y), torch.sin(y / 2), torch.zeros_like(y)), dim=1) | |
| rz = torch.stack((torch.cos(z / 2), torch.zeros_like(z), torch.zeros_like(z), torch.sin(z / 2)), dim=1) | |
| result = None | |
| for coord in order: | |
| if coord == 'x': | |
| r = rx | |
| elif coord == 'y': | |
| r = ry | |
| elif coord == 'z': | |
| r = rz | |
| else: | |
| raise | |
| if result is None: | |
| result = r | |
| else: | |
| result = qmul(result, r) | |
| # Reverse antipodal representation to have a non-negative "w" | |
| if order in ['xyz', 'yzx', 'zxy']: | |
| result *= -1 | |
| return result.view(original_shape) | |
| def expmap_to_quaternion(e): | |
| """ | |
| Convert axis-angle rotations (aka exponential maps) to quaternions. | |
| Stable formula from "Practical Parameterization of Rotations Using the Exponential Map". | |
| Expects a tensor of shape (*, 3), where * denotes any number of dimensions. | |
| Returns a tensor of shape (*, 4). | |
| """ | |
| assert e.shape[-1] == 3 | |
| original_shape = list(e.shape) | |
| original_shape[-1] = 4 | |
| e = e.reshape(-1, 3) | |
| theta = np.linalg.norm(e, axis=1).reshape(-1, 1) | |
| w = np.cos(0.5 * theta).reshape(-1, 1) | |
| xyz = 0.5 * np.sinc(0.5 * theta / np.pi) * e | |
| return np.concatenate((w, xyz), axis=1).reshape(original_shape) | |
| def euler_to_quaternion(e, order): | |
| """ | |
| Convert Euler angles to quaternions. | |
| """ | |
| assert e.shape[-1] == 3 | |
| original_shape = list(e.shape) | |
| original_shape[-1] = 4 | |
| e = e.reshape(-1, 3) | |
| x = e[:, 0] | |
| y = e[:, 1] | |
| z = e[:, 2] | |
| rx = np.stack((np.cos(x / 2), np.sin(x / 2), np.zeros_like(x), np.zeros_like(x)), axis=1) | |
| ry = np.stack((np.cos(y / 2), np.zeros_like(y), np.sin(y / 2), np.zeros_like(y)), axis=1) | |
| rz = np.stack((np.cos(z / 2), np.zeros_like(z), np.zeros_like(z), np.sin(z / 2)), axis=1) | |
| result = None | |
| for coord in order: | |
| if coord == 'x': | |
| r = rx | |
| elif coord == 'y': | |
| r = ry | |
| elif coord == 'z': | |
| r = rz | |
| else: | |
| raise | |
| if result is None: | |
| result = r | |
| else: | |
| result = qmul_np(result, r) | |
| # Reverse antipodal representation to have a non-negative "w" | |
| if order in ['xyz', 'yzx', 'zxy']: | |
| result *= -1 | |
| return result.reshape(original_shape) | |
| def quaternion_to_matrix(quaternions): | |
| """ | |
| Convert rotations given as quaternions to rotation matrices. | |
| Args: | |
| quaternions: quaternions with real part first, | |
| as tensor of shape (..., 4). | |
| Returns: | |
| Rotation matrices as tensor of shape (..., 3, 3). | |
| """ | |
| r, i, j, k = torch.unbind(quaternions, -1) | |
| two_s = 2.0 / (quaternions * quaternions).sum(-1) | |
| o = torch.stack( | |
| ( | |
| 1 - two_s * (j * j + k * k), | |
| two_s * (i * j - k * r), | |
| two_s * (i * k + j * r), | |
| two_s * (i * j + k * r), | |
| 1 - two_s * (i * i + k * k), | |
| two_s * (j * k - i * r), | |
| two_s * (i * k - j * r), | |
| two_s * (j * k + i * r), | |
| 1 - two_s * (i * i + j * j), | |
| ), | |
| -1, | |
| ) | |
| return o.reshape(quaternions.shape[:-1] + (3, 3)) | |
| def quaternion_to_matrix_np(quaternions): | |
| q = torch.from_numpy(quaternions).contiguous().float() | |
| return quaternion_to_matrix(q).numpy() | |
| def quaternion_to_cont6d_np(quaternions): | |
| rotation_mat = quaternion_to_matrix_np(quaternions) | |
| cont_6d = np.concatenate([rotation_mat[..., 0], rotation_mat[..., 1]], axis=-1) | |
| return cont_6d | |
| def quaternion_to_cont6d(quaternions): | |
| rotation_mat = quaternion_to_matrix(quaternions) | |
| cont_6d = torch.cat([rotation_mat[..., 0], rotation_mat[..., 1]], dim=-1) | |
| return cont_6d | |
| def cont6d_to_matrix(cont6d): | |
| assert cont6d.shape[-1] == 6, "The last dimension must be 6" | |
| x_raw = cont6d[..., 0:3] | |
| y_raw = cont6d[..., 3:6] | |
| x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True) | |
| z = torch.cross(x, y_raw, dim=-1) | |
| z = z / torch.norm(z, dim=-1, keepdim=True) | |
| y = torch.cross(z, x, dim=-1) | |
| x = x[..., None] | |
| y = y[..., None] | |
| z = z[..., None] | |
| mat = torch.cat([x, y, z], dim=-1) | |
| return mat | |
| def cont6d_to_matrix_np(cont6d): | |
| q = torch.from_numpy(cont6d).contiguous().float() | |
| return cont6d_to_matrix(q).numpy() | |
| def qpow(q0, t, dtype=torch.float): | |
| ''' q0 : tensor of quaternions | |
| t: tensor of powers | |
| ''' | |
| q0 = qnormalize(q0) | |
| theta0 = torch.acos(q0[..., 0]) | |
| ## if theta0 is close to zero, add epsilon to avoid NaNs | |
| mask = (theta0 <= 10e-10) * (theta0 >= -10e-10) | |
| theta0 = (1 - mask) * theta0 + mask * 10e-10 | |
| v0 = q0[..., 1:] / torch.sin(theta0).view(-1, 1) | |
| if isinstance(t, torch.Tensor): | |
| q = torch.zeros(t.shape + q0.shape) | |
| theta = t.view(-1, 1) * theta0.view(1, -1) | |
| else: ## if t is a number | |
| q = torch.zeros(q0.shape) | |
| theta = t * theta0 | |
| q[..., 0] = torch.cos(theta) | |
| q[..., 1:] = v0 * torch.sin(theta).unsqueeze(-1) | |
| return q.to(dtype) | |
| def qslerp(q0, q1, t): | |
| ''' | |
| q0: starting quaternion | |
| q1: ending quaternion | |
| t: array of points along the way | |
| Returns: | |
| Tensor of Slerps: t.shape + q0.shape | |
| ''' | |
| q0 = qnormalize(q0) | |
| q1 = qnormalize(q1) | |
| q_ = qpow(qmul(q1, qinv(q0)), t) | |
| return qmul(q_, | |
| q0.contiguous().view(torch.Size([1] * len(t.shape)) + q0.shape).expand(t.shape + q0.shape).contiguous()) | |
| def qbetween(v0, v1): | |
| ''' | |
| find the quaternion used to rotate v0 to v1 | |
| ''' | |
| assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)' | |
| assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)' | |
| v = torch.cross(v0, v1) | |
| w = torch.sqrt((v0 ** 2).sum(dim=-1, keepdim=True) * (v1 ** 2).sum(dim=-1, keepdim=True)) + (v0 * v1).sum(dim=-1, | |
| keepdim=True) | |
| return qnormalize(torch.cat([w, v], dim=-1)) | |
| def qbetween_np(v0, v1): | |
| ''' | |
| find the quaternion used to rotate v0 to v1 | |
| ''' | |
| assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)' | |
| assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)' | |
| v0 = torch.from_numpy(v0).float() | |
| v1 = torch.from_numpy(v1).float() | |
| return qbetween(v0, v1).numpy() | |
| def lerp(p0, p1, t): | |
| if not isinstance(t, torch.Tensor): | |
| t = torch.Tensor([t]) | |
| new_shape = t.shape + p0.shape | |
| new_view_t = t.shape + torch.Size([1] * len(p0.shape)) | |
| new_view_p = torch.Size([1] * len(t.shape)) + p0.shape | |
| p0 = p0.view(new_view_p).expand(new_shape) | |
| p1 = p1.view(new_view_p).expand(new_shape) | |
| t = t.view(new_view_t).expand(new_shape) | |
| return p0 + t * (p1 - p0) | |