Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,68 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def rot6d_to_rotmat(x):
|
||||
"""Convert 6D rotation representation to 3x3 rotation matrix.
|
||||
|
||||
Based on Zhou et al., "On the Continuity of Rotation
|
||||
Representations in Neural Networks", CVPR 2019
|
||||
Input:
|
||||
(B,6) Batch of 6-D rotation representations
|
||||
Output:
|
||||
(B,3,3) Batch of corresponding rotation matrices
|
||||
"""
|
||||
x = x.view(-1, 3, 2)
|
||||
a1 = x[:, :, 0]
|
||||
a2 = x[:, :, 1]
|
||||
b1 = F.normalize(a1)
|
||||
b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
|
||||
b3 = torch.cross(b1, b2)
|
||||
return torch.stack((b1, b2, b3), dim=-1)
|
||||
|
||||
|
||||
def batch_rodrigues(theta):
|
||||
"""Convert axis-angle representation to rotation matrix.
|
||||
Args:
|
||||
theta: size = [B, 3]
|
||||
Returns:
|
||||
Rotation matrix corresponding to the quaternion
|
||||
-- size = [B, 3, 3]
|
||||
"""
|
||||
l2norm = torch.norm(theta + 1e-8, p=2, dim=1)
|
||||
angle = torch.unsqueeze(l2norm, -1)
|
||||
normalized = torch.div(theta, angle)
|
||||
angle = angle * 0.5
|
||||
v_cos = torch.cos(angle)
|
||||
v_sin = torch.sin(angle)
|
||||
quat = torch.cat([v_cos, v_sin * normalized], dim=1)
|
||||
return quat_to_rotmat(quat)
|
||||
|
||||
|
||||
def quat_to_rotmat(quat):
|
||||
"""Convert quaternion coefficients to rotation matrix.
|
||||
Args:
|
||||
quat: size = [B, 4] 4 <===>(w, x, y, z)
|
||||
Returns:
|
||||
Rotation matrix corresponding to the quaternion
|
||||
-- size = [B, 3, 3]
|
||||
"""
|
||||
norm_quat = quat
|
||||
norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True)
|
||||
w, x, y, z = norm_quat[:, 0], norm_quat[:, 1],\
|
||||
norm_quat[:, 2], norm_quat[:, 3]
|
||||
|
||||
B = quat.size(0)
|
||||
|
||||
w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
|
||||
wx, wy, wz = w * x, w * y, w * z
|
||||
xy, xz, yz = x * y, x * z, y * z
|
||||
|
||||
rotMat = torch.stack([
|
||||
w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy,
|
||||
w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz,
|
||||
w2 - x2 - y2 + z2
|
||||
],
|
||||
dim=1).view(B, 3, 3)
|
||||
return rotMat
|
||||
Reference in New Issue
Block a user