import torch import numpy as np _EPS4 = np.finfo(float).eps * 4.0 _FLOAT_EPS = np.finfo(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 rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor: """ Converts 6D rotation representation by Zhou et al. [1] to rotation matrix using Gram--Schmidt orthogonalization per Section B of [1]. Args: d6: 6D rotation representation, of size (*, 6) Returns: batch of rotation matrices of size (*, 3, 3) [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. On the Continuity of Rotation Representations in Neural Networks. IEEE Conference on Computer Vision and Pattern Recognition, 2019. Retrieved from http://arxiv.org/abs/1812.07035 """ a1, a2 = d6[..., :3], d6[..., 3:] b1 = F.normalize(a1, dim=-1) b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 b2 = F.normalize(b2, dim=-1) b3 = torch.cross(b1, b2, dim=-1) return torch.stack((b1, b2, b3), dim=-2) # def cont6d_to_matrix_np(cont6d): # q = torch.from_numpy(cont6d).contiguous() # 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) joint_idx = (0,1,2,4,5,7,8,12,16,17,18,19,20,21,60,61,62,63,64,65,59,58,57,56,55, # body joints 37,38,39,66,25,26,27,67,28,29,30,68,34,35,36,69,31,32,33,70, # left hand joints 52,53,54,71,40,41,42,72,43,44,45,73,49,50,51,74,46,47,48,75, # right hand joints 22,15, # jaw, head 57,56, # eyeballs 76,77,78,79,80,81,82,83,84,85, # eyebrow 86,87,88,89, # nose 90,91,92,93,94, # below nose 95,96,97,98,99,100,101,102,103,104,105,106, # eyes 107, # right mouth 108,109,110,111,112, # upper mouth 113, # left mouth 114,115,116,117,118, # lower mouth 119, # right lip 120,121,122, # upper lip 123, # left lip 124,125,126, # lower lip 127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143 # face contour ) def face_z_transform(positions, global_orient, trans): ''' positions: [num_frame, num_joints, 3] global_orient: [num_frame, 3] ''' joints_name = \ ('Pelvis', 'L_Hip', 'R_Hip', 'L_Knee', 'R_Knee', 'L_Ankle', 'R_Ankle', 'Neck', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', 'L_Ear', 'R_Ear', 'L_Eye', 'R_Eye', 'Nose', # body joints 'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', 'L_Thumb_4', 'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Index_4', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Middle_4', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Ring_4', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Pinky_4', # left hand joints 'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', 'R_Thumb_4', 'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Index_4', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Middle_4', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Ring_4', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Pinky_4', # right hand joints *['Face_' + str(i) for i in range(1, 73)] # face keypoints (too many keypoints... omit real names. have same name of keypoints defined in FLAME class) ) root_pos_init = positions[0] assert root_pos_init.shape[0]==len(joints_name) '''All initially face Z+''' r_hip, l_hip, sdr_r, sdr_l = joints_name.index('R_Hip'), joints_name.index('L_Hip'), joints_name.index('R_Shoulder'), joints_name.index('L_Shoulder') across1 = root_pos_init[r_hip] - root_pos_init[l_hip] across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l] across = across1 + across2 across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis] # forward (3,), rotate around y-axis forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1) forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis] target = np.array([[0, 0, 1]]) root_quat_init = qbetween_np(forward_init, target) root_quat_init = np.ones(global_orient.shape[:-1] + (4,)) * root_quat_init root_quat_init = torch.tensor(root_quat_init, dtype=torch.float32).float().cuda() root_matrix_init = quaternion_to_matrix(root_quat_init) global_orient_matrix = axis_angle_to_matrix(global_orient) global_orient_matrix = torch.matmul(root_matrix_init, global_orient_matrix) global_orient = matrix_to_axis_angle(global_orient_matrix) trans = trans.cpu().numpy() '''Put on Floor''' floor_height = positions.min(axis=0).min(axis=0)[1] trans[:, 1] -= floor_height '''XZ at origin''' root_pos_init = positions[0] root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1]) trans = trans - root_pose_init_xz '''All initially face Z+''' trans = torch.from_numpy(trans).float().cuda() trans = qrot(root_quat_init, trans) return global_orient, trans # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. # Check PYTORCH3D_LICENCE before use import functools from typing import Optional import torch import torch.nn.functional as F """ The transformation matrices returned from the functions in this file assume the points on which the transformation will be applied are column vectors. i.e. the R matrix is structured as R = [ [Rxx, Rxy, Rxz], [Ryx, Ryy, Ryz], [Rzx, Rzy, Rzz], ] # (3, 3) This matrix can be applied to column vectors by post multiplication by the points e.g. points = [[0], [1], [2]] # (3 x 1) xyz coordinates of a point transformed_points = R * points To apply the same matrix to points which are row vectors, the R matrix can be transposed and pre multiplied by the points: e.g. points = [[0, 1, 2]] # (1 x 3) xyz coordinates of a point transformed_points = points * R.transpose(1, 0) """ 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 _copysign(a, b): """ Return a tensor where each element has the absolute value taken from the, corresponding element of a, with sign taken from the corresponding element of b. This is like the standard copysign floating-point operation, but is not careful about negative 0 and NaN. Args: a: source tensor. b: tensor whose signs will be used, of the same shape as a. Returns: Tensor of the same shape as a with the signs of b. """ signs_differ = (a < 0) != (b < 0) return torch.where(signs_differ, -a, a) def _sqrt_positive_part(x): """ Returns torch.sqrt(torch.max(0, x)) but with a zero subgradient where x is 0. """ ret = torch.zeros_like(x) positive_mask = x > 0 ret[positive_mask] = torch.sqrt(x[positive_mask]) return ret def matrix_to_quaternion(matrix): """ Convert rotations given as rotation matrices to quaternions. Args: matrix: Rotation matrices as tensor of shape (..., 3, 3). Returns: quaternions with real part first, as tensor of shape (..., 4). """ if matrix.size(-1) != 3 or matrix.size(-2) != 3: raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.") m00 = matrix[..., 0, 0] m11 = matrix[..., 1, 1] m22 = matrix[..., 2, 2] o0 = 0.5 * _sqrt_positive_part(1 + m00 + m11 + m22) x = 0.5 * _sqrt_positive_part(1 + m00 - m11 - m22) y = 0.5 * _sqrt_positive_part(1 - m00 + m11 - m22) z = 0.5 * _sqrt_positive_part(1 - m00 - m11 + m22) o1 = _copysign(x, matrix[..., 2, 1] - matrix[..., 1, 2]) o2 = _copysign(y, matrix[..., 0, 2] - matrix[..., 2, 0]) o3 = _copysign(z, matrix[..., 1, 0] - matrix[..., 0, 1]) return torch.stack((o0, o1, o2, o3), -1) def _axis_angle_rotation(axis: str, angle): """ Return the rotation matrices for one of the rotations about an axis of which Euler angles describe, for each value of the angle given. Args: axis: Axis label "X" or "Y or "Z". angle: any shape tensor of Euler angles in radians Returns: Rotation matrices as tensor of shape (..., 3, 3). """ cos = torch.cos(angle) sin = torch.sin(angle) one = torch.ones_like(angle) zero = torch.zeros_like(angle) if axis == "X": R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos) if axis == "Y": R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos) if axis == "Z": R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one) return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3)) def euler_angles_to_matrix(euler_angles, convention: str): """ Convert rotations given as Euler angles in radians to rotation matrices. Args: euler_angles: Euler angles in radians as tensor of shape (..., 3). convention: Convention string of three uppercase letters from {"X", "Y", and "Z"}. Returns: Rotation matrices as tensor of shape (..., 3, 3). """ if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3: raise ValueError("Invalid input euler angles.") if len(convention) != 3: raise ValueError("Convention must have 3 letters.") if convention[1] in (convention[0], convention[2]): raise ValueError(f"Invalid convention {convention}.") for letter in convention: if letter not in ("X", "Y", "Z"): raise ValueError(f"Invalid letter {letter} in convention string.") matrices = map(_axis_angle_rotation, convention, torch.unbind(euler_angles, -1)) return functools.reduce(torch.matmul, matrices) def _angle_from_tan( axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool ): """ Extract the first or third Euler angle from the two members of the matrix which are positive constant times its sine and cosine. Args: axis: Axis label "X" or "Y or "Z" for the angle we are finding. other_axis: Axis label "X" or "Y or "Z" for the middle axis in the convention. data: Rotation matrices as tensor of shape (..., 3, 3). horizontal: Whether we are looking for the angle for the third axis, which means the relevant entries are in the same row of the rotation matrix. If not, they are in the same column. tait_bryan: Whether the first and third axes in the convention differ. Returns: Euler Angles in radians for each matrix in data as a tensor of shape (...). """ i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis] if horizontal: i2, i1 = i1, i2 even = (axis + other_axis) in ["XY", "YZ", "ZX"] if horizontal == even: return torch.atan2(data[..., i1], data[..., i2]) if tait_bryan: return torch.atan2(-data[..., i2], data[..., i1]) return torch.atan2(data[..., i2], -data[..., i1]) def _index_from_letter(letter: str): if letter == "X": return 0 if letter == "Y": return 1 if letter == "Z": return 2 def matrix_to_euler_angles(matrix, convention: str): """ Convert rotations given as rotation matrices to Euler angles in radians. Args: matrix: Rotation matrices as tensor of shape (..., 3, 3). convention: Convention string of three uppercase letters. Returns: Euler angles in radians as tensor of shape (..., 3). """ if len(convention) != 3: raise ValueError("Convention must have 3 letters.") if convention[1] in (convention[0], convention[2]): raise ValueError(f"Invalid convention {convention}.") for letter in convention: if letter not in ("X", "Y", "Z"): raise ValueError(f"Invalid letter {letter} in convention string.") if matrix.size(-1) != 3 or matrix.size(-2) != 3: raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.") i0 = _index_from_letter(convention[0]) i2 = _index_from_letter(convention[2]) tait_bryan = i0 != i2 if tait_bryan: central_angle = torch.asin( matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0) ) else: central_angle = torch.acos(matrix[..., i0, i0]) o = ( _angle_from_tan( convention[0], convention[1], matrix[..., i2], False, tait_bryan ), central_angle, _angle_from_tan( convention[2], convention[1], matrix[..., i0, :], True, tait_bryan ), ) return torch.stack(o, -1) def random_quaternions( n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False ): """ Generate random quaternions representing rotations, i.e. versors with nonnegative real part. Args: n: Number of quaternions in a batch to return. dtype: Type to return. device: Desired device of returned tensor. Default: uses the current device for the default tensor type. requires_grad: Whether the resulting tensor should have the gradient flag set. Returns: Quaternions as tensor of shape (N, 4). """ o = torch.randn((n, 4), dtype=dtype, device=device, requires_grad=requires_grad) s = (o * o).sum(1) o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None] return o def random_rotations( n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False ): """ Generate random rotations as 3x3 rotation matrices. Args: n: Number of rotation matrices in a batch to return. dtype: Type to return. device: Device of returned tensor. Default: if None, uses the current device for the default tensor type. requires_grad: Whether the resulting tensor should have the gradient flag set. Returns: Rotation matrices as tensor of shape (n, 3, 3). """ quaternions = random_quaternions( n, dtype=dtype, device=device, requires_grad=requires_grad ) return quaternion_to_matrix(quaternions) def random_rotation( dtype: Optional[torch.dtype] = None, device=None, requires_grad=False ): """ Generate a single random 3x3 rotation matrix. Args: dtype: Type to return device: Device of returned tensor. Default: if None, uses the current device for the default tensor type requires_grad: Whether the resulting tensor should have the gradient flag set Returns: Rotation matrix as tensor of shape (3, 3). """ return random_rotations(1, dtype, device, requires_grad)[0] def standardize_quaternion(quaternions): """ Convert a unit quaternion to a standard form: one in which the real part is non negative. Args: quaternions: Quaternions with real part first, as tensor of shape (..., 4). Returns: Standardized quaternions as tensor of shape (..., 4). """ return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions) def quaternion_raw_multiply(a, b): """ Multiply two quaternions. Usual torch rules for broadcasting apply. Args: a: Quaternions as tensor of shape (..., 4), real part first. b: Quaternions as tensor of shape (..., 4), real part first. Returns: The product of a and b, a tensor of quaternions shape (..., 4). """ aw, ax, ay, az = torch.unbind(a, -1) bw, bx, by, bz = torch.unbind(b, -1) ow = aw * bw - ax * bx - ay * by - az * bz ox = aw * bx + ax * bw + ay * bz - az * by oy = aw * by - ax * bz + ay * bw + az * bx oz = aw * bz + ax * by - ay * bx + az * bw return torch.stack((ow, ox, oy, oz), -1) def quaternion_multiply(a, b): """ Multiply two quaternions representing rotations, returning the quaternion representing their composition, i.e. the versor with nonnegative real part. Usual torch rules for broadcasting apply. Args: a: Quaternions as tensor of shape (..., 4), real part first. b: Quaternions as tensor of shape (..., 4), real part first. Returns: The product of a and b, a tensor of quaternions of shape (..., 4). """ ab = quaternion_raw_multiply(a, b) return standardize_quaternion(ab) def quaternion_invert(quaternion): """ Given a quaternion representing rotation, get the quaternion representing its inverse. Args: quaternion: Quaternions as tensor of shape (..., 4), with real part first, which must be versors (unit quaternions). Returns: The inverse, a tensor of quaternions of shape (..., 4). """ return quaternion * quaternion.new_tensor([1, -1, -1, -1]) def quaternion_apply(quaternion, point): """ Apply the rotation given by a quaternion to a 3D point. Usual torch rules for broadcasting apply. Args: quaternion: Tensor of quaternions, real part first, of shape (..., 4). point: Tensor of 3D points of shape (..., 3). Returns: Tensor of rotated points of shape (..., 3). """ if point.size(-1) != 3: raise ValueError(f"Points are not in 3D, f{point.shape}.") real_parts = point.new_zeros(point.shape[:-1] + (1,)) point_as_quaternion = torch.cat((real_parts, point), -1) out = quaternion_raw_multiply( quaternion_raw_multiply(quaternion, point_as_quaternion), quaternion_invert(quaternion), ) return out[..., 1:] def axis_angle_to_matrix(axis_angle): """ Convert rotations given as axis/angle to rotation matrices. Args: axis_angle: Rotations given as a vector in axis angle form, as a tensor of shape (..., 3), where the magnitude is the angle turned anticlockwise in radians around the vector's direction. Returns: Rotation matrices as tensor of shape (..., 3, 3). """ return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle)) def matrix_to_axis_angle(matrix): """ Convert rotations given as rotation matrices to axis/angle. Args: matrix: Rotation matrices as tensor of shape (..., 3, 3). Returns: Rotations given as a vector in axis angle form, as a tensor of shape (..., 3), where the magnitude is the angle turned anticlockwise in radians around the vector's direction. """ return quaternion_to_axis_angle(matrix_to_quaternion(matrix)) def axis_angle_to_quaternion(axis_angle): """ Convert rotations given as axis/angle to quaternions. Args: axis_angle: Rotations given as a vector in axis angle form, as a tensor of shape (..., 3), where the magnitude is the angle turned anticlockwise in radians around the vector's direction. Returns: quaternions with real part first, as tensor of shape (..., 4). """ angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True) half_angles = 0.5 * angles eps = 1e-6 small_angles = angles.abs() < eps sin_half_angles_over_angles = torch.empty_like(angles) sin_half_angles_over_angles[~small_angles] = ( torch.sin(half_angles[~small_angles]) / angles[~small_angles] ) # for x small, sin(x/2) is about x/2 - (x/2)^3/6 # so sin(x/2)/x is about 1/2 - (x*x)/48 sin_half_angles_over_angles[small_angles] = ( 0.5 - (angles[small_angles] * angles[small_angles]) / 48 ) quaternions = torch.cat( [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1 ) return quaternions def quaternion_to_axis_angle(quaternions): """ Convert rotations given as quaternions to axis/angle. Args: quaternions: quaternions with real part first, as tensor of shape (..., 4). Returns: Rotations given as a vector in axis angle form, as a tensor of shape (..., 3), where the magnitude is the angle turned anticlockwise in radians around the vector's direction. """ norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True) half_angles = torch.atan2(norms, quaternions[..., :1]) angles = 2 * half_angles eps = 1e-6 small_angles = angles.abs() < eps sin_half_angles_over_angles = torch.empty_like(angles) sin_half_angles_over_angles[~small_angles] = ( torch.sin(half_angles[~small_angles]) / angles[~small_angles] ) # for x small, sin(x/2) is about x/2 - (x/2)^3/6 # so sin(x/2)/x is about 1/2 - (x*x)/48 sin_half_angles_over_angles[small_angles] = ( 0.5 - (angles[small_angles] * angles[small_angles]) / 48 ) return quaternions[..., 1:] / sin_half_angles_over_angles def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor: """ Converts 6D rotation representation by Zhou et al. [1] to rotation matrix using Gram--Schmidt orthogonalisation per Section B of [1]. Args: d6: 6D rotation representation, of size (*, 6) Returns: batch of rotation matrices of size (*, 3, 3) [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. On the Continuity of Rotation Representations in Neural Networks. IEEE Conference on Computer Vision and Pattern Recognition, 2019. Retrieved from http://arxiv.org/abs/1812.07035 """ a1, a2 = d6[..., :3], d6[..., 3:] b1 = F.normalize(a1, dim=-1) b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 b2 = F.normalize(b2, dim=-1) b3 = torch.cross(b1, b2, dim=-1) return torch.stack((b1, b2, b3), dim=-2) def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor: """ Converts rotation matrices to 6D rotation representation by Zhou et al. [1] by dropping the last row. Note that 6D representation is not unique. Args: matrix: batch of rotation matrices of size (*, 3, 3) Returns: 6D rotation representation, of size (*, 6) [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. On the Continuity of Rotation Representations in Neural Networks. IEEE Conference on Computer Vision and Pattern Recognition, 2019. Retrieved from http://arxiv.org/abs/1812.07035 """ return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6) def canonicalize_smplh(poses, trans=None): bs, nframes, njoints = poses.shape[:3] global_orient = poses[:, :, 0] # first global rotations rot2d = matrix_to_axis_angle(global_orient[:, 0]) # rot2d[:, :2] = 0 # Remove the rotation along the vertical axis rot2d = axis_angle_to_matrix(rot2d) # Rotate the global rotation to eliminate Z rotations global_orient = torch.einsum("ikj,imkl->imjl", rot2d, global_orient) # Construct canonicalized version of x xc = torch.cat((global_orient[:, :, None], poses[:, :, 1:]), dim=2) if trans is not None: vel = trans[:, 1:] - trans[:, :-1] # Turn the translation as well vel = torch.einsum("ikj,ilk->ilj", rot2d, vel) trans = torch.cat((torch.zeros(bs, 1, 3, device=vel.device), torch.cumsum(vel, 1)), 1) return xc, trans else: return xc # Added def matrix_of_angles(cos, sin, inv=False, dim=2): assert dim in [2, 3] sin = -sin if inv else sin if dim == 2: row1 = torch.stack((cos, -sin), axis=-1) row2 = torch.stack((sin, cos), axis=-1) return torch.stack((row1, row2), axis=-2) elif dim == 3: row1 = torch.stack((cos, -sin, 0 * cos), axis=-1) row2 = torch.stack((sin, cos, 0 * cos), axis=-1) row3 = torch.stack((0 * sin, 0 * cos, 1 + 0 * cos), axis=-1) return torch.stack((row1, row2, row3), axis=-2)