OpenMotionLab/MotionGPT
118
1# -*- coding: utf-8 -*-2 3# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is4# holder of all proprietary rights on this computer program.5# You can only use this computer program if you have closed6# a license agreement with MPG or you get the right to use the computer7# program from someone who is authorized to grant you that right.8# Any use of the computer program without a valid license is prohibited and9# liable to prosecution.10#11# Copyright©2019 Max-Planck-Gesellschaft zur Förderung12# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute13# for Intelligent Systems. All rights reserved.14#15# Contact: ps-license@tuebingen.mpg.de16 17import torch18import numpy as np19from torch.nn import functional as F20 21 22def axis_angle_to_quaternion(axis_angle):23 """24 Convert rotations given as axis/angle to quaternions.25 26 Args:27 axis_angle: Rotations given as a vector in axis angle form,28 as a tensor of shape (..., 3), where the magnitude is29 the angle turned anticlockwise in radians around the30 vector's direction.31 32 Returns:33 quaternions with real part first, as tensor of shape (..., 4).34 """35 angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)36 half_angles = 0.5 * angles37 eps = 1e-638 small_angles = angles.abs() < eps39 sin_half_angles_over_angles = torch.empty_like(angles)40 sin_half_angles_over_angles[~small_angles] = (41 torch.sin(half_angles[~small_angles]) / angles[~small_angles])42 # for x small, sin(x/2) is about x/2 - (x/2)^3/643 # so sin(x/2)/x is about 1/2 - (x*x)/4844 sin_half_angles_over_angles[small_angles] = (45 0.5 - (angles[small_angles] * angles[small_angles]) / 48)46 quaternions = torch.cat(47 [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles],48 dim=-1)49 return quaternions50 51 52def quaternion_to_matrix(quaternions):53 """54 Convert rotations given as quaternions to rotation matrices.55 56 Args:57 quaternions: quaternions with real part first,58 as tensor of shape (..., 4).59 60 Returns:61 Rotation matrices as tensor of shape (..., 3, 3).62 """63 r, i, j, k = torch.unbind(quaternions, -1)64 two_s = 2.0 / (quaternions * quaternions).sum(-1)65 66 o = torch.stack(67 (68 1 - two_s * (j * j + k * k),69 two_s * (i * j - k * r),70 two_s * (i * k + j * r),71 two_s * (i * j + k * r),72 1 - two_s * (i * i + k * k),73 two_s * (j * k - i * r),74 two_s * (i * k - j * r),75 two_s * (j * k + i * r),76 1 - two_s * (i * i + j * j),77 ),78 -1,79 )80 return o.reshape(quaternions.shape[:-1] + (3, 3))81 82 83def axis_angle_to_matrix(axis_angle):84 """85 Convert rotations given as axis/angle to rotation matrices.86 87 Args:88 axis_angle: Rotations given as a vector in axis angle form,89 as a tensor of shape (..., 3), where the magnitude is90 the angle turned anticlockwise in radians around the91 vector's direction.92 93 Returns:94 Rotation matrices as tensor of shape (..., 3, 3).95 """96 return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))97 98 99def matrix_of_angles(cos, sin, inv=False, dim=2):100 assert dim in [2, 3]101 sin = -sin if inv else sin102 if dim == 2:103 row1 = torch.stack((cos, -sin), axis=-1)104 row2 = torch.stack((sin, cos), axis=-1)105 return torch.stack((row1, row2), axis=-2)106 elif dim == 3:107 row1 = torch.stack((cos, -sin, 0 * cos), axis=-1)108 row2 = torch.stack((sin, cos, 0 * cos), axis=-1)109 row3 = torch.stack((0 * sin, 0 * cos, 1 + 0 * cos), axis=-1)110 return torch.stack((row1, row2, row3), axis=-2)111 112 113def matrot2axisangle(matrots):114 # This function is borrowed from https://github.com/davrempe/humor/utils/transforms.py115 # axisang N x 3116 '''117 :param matrots: N*num_joints*9118 :return: N*num_joints*3119 '''120 import cv2121 batch_size = matrots.shape[0]122 matrots = matrots.reshape([batch_size, -1, 9])123 out_axisangle = []124 for mIdx in range(matrots.shape[0]):125 cur_axisangle = []126 for jIdx in range(matrots.shape[1]):127 a = cv2.Rodrigues(matrots[mIdx,128 jIdx:jIdx + 1, :].reshape(3,129 3))[0].reshape(130 (1, 3))131 cur_axisangle.append(a)132 133 out_axisangle.append(np.array(cur_axisangle).reshape([1, -1, 3]))134 return np.vstack(out_axisangle)135 136 137def axisangle2matrots(axisangle):138 # This function is borrowed from https://github.com/davrempe/humor/utils/transforms.py139 # axisang N x 3140 '''141 :param axisangle: N*num_joints*3142 :return: N*num_joints*9143 '''144 import cv2145 batch_size = axisangle.shape[0]146 axisangle = axisangle.reshape([batch_size, -1, 3])147 out_matrot = []148 for mIdx in range(axisangle.shape[0]):149 cur_axisangle = []150 for jIdx in range(axisangle.shape[1]):151 a = cv2.Rodrigues(axisangle[mIdx, jIdx:jIdx + 1, :].reshape(1,152 3))[0]153 cur_axisangle.append(a)154 155 out_matrot.append(np.array(cur_axisangle).reshape([1, -1, 9]))156 return np.vstack(out_matrot)157 158 159def batch_rodrigues(axisang):160 # This function is borrowed from https://github.com/MandyMo/pytorch_HMR/blob/master/src/util.py#L37161 # axisang N x 3162 axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1)163 angle = torch.unsqueeze(axisang_norm, -1)164 axisang_normalized = torch.div(axisang, angle)165 angle = angle * 0.5166 v_cos = torch.cos(angle)167 v_sin = torch.sin(angle)168 169 quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)170 rot_mat = quat2mat(quat)171 rot_mat = rot_mat.view(rot_mat.shape[0], 9)172 return rot_mat173 174 175def quat2mat(quat):176 """177 This function is borrowed from https://github.com/MandyMo/pytorch_HMR/blob/master/src/util.py#L50178 179 Convert quaternion coefficients to rotation matrix.180 Args:181 quat: size = [batch_size, 4] 4 <===>(w, x, y, z)182 Returns:183 Rotation matrix corresponding to the quaternion -- size = [batch_size, 3, 3]184 """185 norm_quat = quat186 norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True)187 w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:,188 2], norm_quat[:,189 3]190 191 batch_size = quat.size(0)192 193 w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)194 wx, wy, wz = w * x, w * y, w * z195 xy, xz, yz = x * y, x * z, y * z196 197 rotMat = torch.stack([198 w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy,199 w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz,200 w2 - x2 - y2 + z2201 ],202 dim=1).view(batch_size, 3, 3)203 return rotMat204 205 206def rotation_matrix_to_angle_axis(rotation_matrix):207 """208 This function is borrowed from https://github.com/kornia/kornia209 210 Convert 3x4 rotation matrix to Rodrigues vector211 212 Args:213 rotation_matrix (Tensor): rotation matrix.214 215 Returns:216 Tensor: Rodrigues vector transformation.217 218 Shape:219 - Input: :math:`(N, 3, 4)`220 - Output: :math:`(N, 3)`221 222 Example:223 >>> input = torch.rand(2, 3, 4) # Nx4x4224 >>> output = tgm.rotation_matrix_to_angle_axis(input) # Nx3225 """226 if rotation_matrix.shape[1:] == (3, 3):227 rot_mat = rotation_matrix.reshape(-1, 3, 3)228 hom = torch.tensor([0, 0, 1],229 dtype=torch.float32,230 device=rotation_matrix.device).reshape(231 1, 3, 1).expand(rot_mat.shape[0], -1, -1)232 rotation_matrix = torch.cat([rot_mat, hom], dim=-1)233 234 quaternion = rotation_matrix_to_quaternion(rotation_matrix)235 aa = quaternion_to_angle_axis(quaternion)236 aa[torch.isnan(aa)] = 0.0237 return aa238 239 240def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:241 """242 This function is borrowed from https://github.com/kornia/kornia243 244 Convert quaternion vector to angle axis of rotation.245 246 Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h247 248 Args:249 quaternion (torch.Tensor): tensor with quaternions.250 251 Return:252 torch.Tensor: tensor with angle axis of rotation.253 254 Shape:255 - Input: :math:`(*, 4)` where `*` means, any number of dimensions256 - Output: :math:`(*, 3)`257 258 Example:259 >>> quaternion = torch.rand(2, 4) # Nx4260 >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion) # Nx3261 """262 if not torch.is_tensor(quaternion):263 raise TypeError("Input type is not a torch.Tensor. Got {}".format(264 type(quaternion)))265 266 if not quaternion.shape[-1] == 4:267 raise ValueError(268 "Input must be a tensor of shape Nx4 or 4. Got {}".format(269 quaternion.shape))270 # unpack input and compute conversion271 q1: torch.Tensor = quaternion[..., 1]272 q2: torch.Tensor = quaternion[..., 2]273 q3: torch.Tensor = quaternion[..., 3]274 sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3275 276 sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)277 cos_theta: torch.Tensor = quaternion[..., 0]278 two_theta: torch.Tensor = 2.0 * torch.where(279 cos_theta < 0.0, torch.atan2(-sin_theta, -cos_theta),280 torch.atan2(sin_theta, cos_theta))281 282 k_pos: torch.Tensor = two_theta / sin_theta283 k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)284 k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)285 286 angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]287 angle_axis[..., 0] += q1 * k288 angle_axis[..., 1] += q2 * k289 angle_axis[..., 2] += q3 * k290 return angle_axis291 292 293def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):294 """295 This function is borrowed from https://github.com/kornia/kornia296 297 Convert 3x4 rotation matrix to 4d quaternion vector298 299 This algorithm is based on algorithm described in300 https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201301 302 Args:303 rotation_matrix (Tensor): the rotation matrix to convert.304 305 Return:306 Tensor: the rotation in quaternion307 308 Shape:309 - Input: :math:`(N, 3, 4)`310 - Output: :math:`(N, 4)`311 312 Example:313 >>> input = torch.rand(4, 3, 4) # Nx3x4314 >>> output = tgm.rotation_matrix_to_quaternion(input) # Nx4315 """316 if not torch.is_tensor(rotation_matrix):317 raise TypeError("Input type is not a torch.Tensor. Got {}".format(318 type(rotation_matrix)))319 320 if len(rotation_matrix.shape) > 3:321 raise ValueError(322 "Input size must be a three dimensional tensor. Got {}".format(323 rotation_matrix.shape))324 if not rotation_matrix.shape[-2:] == (3, 4):325 raise ValueError(326 "Input size must be a N x 3 x 4 tensor. Got {}".format(327 rotation_matrix.shape))328 329 rmat_t = torch.transpose(rotation_matrix, 1, 2)330 331 mask_d2 = rmat_t[:, 2, 2] < eps332 333 mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]334 mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]335 336 t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]337 q0 = torch.stack([338 rmat_t[:, 1, 2] - rmat_t[:, 2, 1], t0,339 rmat_t[:, 0, 1] + rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2]340 ], -1)341 t0_rep = t0.repeat(4, 1).t()342 343 t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]344 q1 = torch.stack([345 rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] + rmat_t[:, 1, 0],346 t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]347 ], -1)348 t1_rep = t1.repeat(4, 1).t()349 350 t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]351 q2 = torch.stack([352 rmat_t[:, 0, 1] - rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2],353 rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2354 ], -1)355 t2_rep = t2.repeat(4, 1).t()356 357 t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]358 q3 = torch.stack([359 t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],360 rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] - rmat_t[:, 1, 0]361 ], -1)362 t3_rep = t3.repeat(4, 1).t()363 364 mask_c0 = mask_d2 * mask_d0_d1365 mask_c1 = mask_d2 * ~mask_d0_d1366 mask_c2 = ~mask_d2 * mask_d0_nd1367 mask_c3 = ~mask_d2 * ~mask_d0_nd1368 mask_c0 = mask_c0.view(-1, 1).type_as(q0)369 mask_c1 = mask_c1.view(-1, 1).type_as(q1)370 mask_c2 = mask_c2.view(-1, 1).type_as(q2)371 mask_c3 = mask_c3.view(-1, 1).type_as(q3)372 373 q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3374 q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 + # noqa375 t2_rep * mask_c2 + t3_rep * mask_c3) # noqa376 q *= 0.5377 return q378 379 380def estimate_translation_np(S,381 joints_2d,382 joints_conf,383 focal_length=5000.,384 img_size=224.):385 """386 This function is borrowed from https://github.com/nkolot/SPIN/utils/geometry.py387 388 Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.389 Input:390 S: (25, 3) 3D joint locations391 joints: (25, 3) 2D joint locations and confidence392 Returns:393 (3,) camera translation vector394 """395 396 num_joints = S.shape[0]397 # focal length398 f = np.array([focal_length, focal_length])399 # optical center400 center = np.array([img_size / 2., img_size / 2.])401 402 # transformations403 Z = np.reshape(np.tile(S[:, 2], (2, 1)).T, -1)404 XY = np.reshape(S[:, 0:2], -1)405 O = np.tile(center, num_joints)406 F = np.tile(f, num_joints)407 weight2 = np.reshape(np.tile(np.sqrt(joints_conf), (2, 1)).T, -1)408 409 # least squares410 Q = np.array([411 F * np.tile(np.array([1, 0]), num_joints),412 F * np.tile(np.array([0, 1]), num_joints),413 O - np.reshape(joints_2d, -1)414 ]).T415 c = (np.reshape(joints_2d, -1) - O) * Z - F * XY416 417 # weighted least squares418 W = np.diagflat(weight2)419 Q = np.dot(W, Q)420 c = np.dot(W, c)421 422 # square matrix423 A = np.dot(Q.T, Q)424 b = np.dot(Q.T, c)425 426 # solution427 trans = np.linalg.solve(A, b)428 429 return trans430 431 432def estimate_translation(S, joints_2d, focal_length=5000., img_size=224.):433 """434 This function is borrowed from https://github.com/nkolot/SPIN/utils/geometry.py435 436 Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.437 Input:438 S: (B, 49, 3) 3D joint locations439 joints: (B, 49, 3) 2D joint locations and confidence440 Returns:441 (B, 3) camera translation vectors442 """443 444 device = S.device445 # Use only joints 25:49 (GT joints)446 S = S[:, 25:, :].cpu().numpy()447 joints_2d = joints_2d[:, 25:, :].cpu().numpy()448 joints_conf = joints_2d[:, :, -1]449 joints_2d = joints_2d[:, :, :-1]450 trans = np.zeros((S.shape[0], 3), dtype=np.float6432)451 # Find the translation for each example in the batch452 for i in range(S.shape[0]):453 S_i = S[i]454 joints_i = joints_2d[i]455 conf_i = joints_conf[i]456 trans[i] = estimate_translation_np(S_i,457 joints_i,458 conf_i,459 focal_length=focal_length,460 img_size=img_size)461 return torch.from_numpy(trans).to(device)462 463 464def rot6d_to_rotmat_spin(x):465 """Convert 6D rotation representation to 3x3 rotation matrix.466 Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019467 Input:468 (B,6) Batch of 6-D rotation representations469 Output:470 (B,3,3) Batch of corresponding rotation matrices471 """472 x = x.view(-1, 3, 2)473 a1 = x[:, :, 0]474 a2 = x[:, :, 1]475 b1 = F.normalize(a1)476 b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)477 478 # inp = a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1479 # denom = inp.pow(2).sum(dim=1).sqrt().unsqueeze(-1) + 1e-8480 # b2 = inp / denom481 482 b3 = torch.cross(b1, b2)483 return torch.stack((b1, b2, b3), dim=-1)484 485 486def rot6d_to_rotmat(x):487 x = x.view(-1, 3, 2)488 489 # Normalize the first vector490 b1 = F.normalize(x[:, :, 0], dim=1, eps=1e-6)491 492 dot_prod = torch.sum(b1 * x[:, :, 1], dim=1, keepdim=True)493 # Compute the second vector by finding the orthogonal complement to it494 b2 = F.normalize(x[:, :, 1] - dot_prod * b1, dim=-1, eps=1e-6)495 496 # Finish building the basis by taking the cross product497 b3 = torch.cross(b1, b2, dim=1)498 rot_mats = torch.stack([b1, b2, b3], dim=-1)499 500 return rot_mats501 502 503import mGPT.utils.rotation_conversions as rotation_conversions504 505 506def rot6d(x_rotations, pose_rep):507 time, njoints, feats = x_rotations.shape508 509 # Compute rotations (convert only masked sequences output)510 if pose_rep == "rotvec":511 rotations = rotation_conversions.axis_angle_to_matrix(x_rotations)512 elif pose_rep == "rotmat":513 rotations = x_rotations.view(njoints, 3, 3)514 elif pose_rep == "rotquat":515 rotations = rotation_conversions.quaternion_to_matrix(x_rotations)516 elif pose_rep == "rot6d":517 rotations = rotation_conversions.rotation_6d_to_matrix(x_rotations)518 else:519 raise NotImplementedError("No geometry for this one.")520 521 rotations_6d = rotation_conversions.matrix_to_rotation_6d(rotations)522 return rotations_6d523 524 525def rot6d_batch(x_rotations, pose_rep):526 nsamples, time, njoints, feats = x_rotations.shape527 528 # Compute rotations (convert only masked sequences output)529 if pose_rep == "rotvec":530 rotations = rotation_conversions.axis_angle_to_matrix(x_rotations)531 elif pose_rep == "rotmat":532 rotations = x_rotations.view(-1, njoints, 3, 3)533 elif pose_rep == "rotquat":534 rotations = rotation_conversions.quaternion_to_matrix(x_rotations)535 elif pose_rep == "rot6d":536 rotations = rotation_conversions.rotation_6d_to_matrix(x_rotations)537 else:538 raise NotImplementedError("No geometry for this one.")539 540 rotations_6d = rotation_conversions.matrix_to_rotation_6d(rotations)541 return rotations_6d542 543 544def rot6d_to_rotvec_batch(pose):545 # nsamples, time, njoints, feats = rot6d.shape546 bs, nfeats = pose.shape547 rot6d = pose.reshape(bs, 24, 6)548 rotations = rotation_conversions.rotation_6d_to_matrix(rot6d)549 rotvec = rotation_conversions.matrix_to_axis_angle(rotations)550 return rotvec.reshape(bs, 24 * 3)551