mirror of
https://gitcode.com/gh_mirrors/eas/EasyFace.git
synced 2026-04-30 16:10:17 +00:00
73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
# This code is borrowed and modified from Actor,
|
|
# made publicly available under MIT license at https://github.com/Mathux/ACTOR
|
|
|
|
import torch
|
|
|
|
|
|
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 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 recover_root_rot_pos(data):
|
|
rot_vel = data[..., 0]
|
|
r_rot_ang = torch.zeros_like(rot_vel).to(data.device)
|
|
'''Get Y-axis rotation from rotation velocity'''
|
|
r_rot_ang[..., 1:] = rot_vel[..., :-1]
|
|
r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)
|
|
|
|
r_rot_quat = torch.zeros(data.shape[:-1] + (4, )).to(data.device)
|
|
r_rot_quat[..., 0] = torch.cos(r_rot_ang)
|
|
r_rot_quat[..., 2] = torch.sin(r_rot_ang)
|
|
|
|
r_pos = torch.zeros(data.shape[:-1] + (3, )).to(data.device)
|
|
r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]
|
|
'''Add Y-axis rotation to root position'''
|
|
r_pos = qrot(qinv(r_rot_quat), r_pos)
|
|
|
|
r_pos = torch.cumsum(r_pos, dim=-2)
|
|
|
|
r_pos[..., 1] = data[..., 3]
|
|
return r_rot_quat, r_pos
|
|
|
|
|
|
def recover_from_ric(data, joints_num):
|
|
r_rot_quat, r_pos = recover_root_rot_pos(data)
|
|
positions = data[..., 4:(joints_num - 1) * 3 + 4]
|
|
positions = positions.view(positions.shape[:-1] + (-1, 3))
|
|
'''Add Y-axis rotation to local joints'''
|
|
positions = qrot(
|
|
qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4, )),
|
|
positions)
|
|
'''Add root XZ to joints'''
|
|
positions[..., 0] += r_pos[..., 0:1]
|
|
positions[..., 2] += r_pos[..., 2:3]
|
|
'''Concate root and joints'''
|
|
positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)
|
|
|
|
return positions
|