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,4 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .smpl import SMPL
|
||||
|
||||
__all__ = ['SMPL']
|
||||
@@ -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
|
||||
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def resize(input,
|
||||
size=None,
|
||||
scale_factor=None,
|
||||
mode='nearest',
|
||||
align_corners=None,
|
||||
warning=True):
|
||||
if warning:
|
||||
if size is not None and align_corners:
|
||||
input_h, input_w = tuple(int(x) for x in input.shape[2:])
|
||||
output_h, output_w = tuple(int(x) for x in size)
|
||||
if output_h > input_h or output_w > output_h:
|
||||
if ((output_h > 1 and output_w > 1 and input_h > 1
|
||||
and input_w > 1) and (output_h - 1) % (input_h - 1)
|
||||
and (output_w - 1) % (input_w - 1)):
|
||||
warnings.warn(
|
||||
f'When align_corners={align_corners}, '
|
||||
'the output would more aligned if '
|
||||
f'input size {(input_h, input_w)} is `x+1` and '
|
||||
f'out size {(output_h, output_w)} is `nx+1`')
|
||||
if isinstance(size, torch.Size):
|
||||
size = tuple(int(x) for x in size)
|
||||
return F.interpolate(input, size, scale_factor, mode, align_corners)
|
||||
@@ -0,0 +1,184 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from ..builder import MESH_MODELS
|
||||
|
||||
try:
|
||||
from smplx import SMPL as SMPL_
|
||||
has_smpl = True
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
has_smpl = False
|
||||
|
||||
|
||||
@MESH_MODELS.register_module()
|
||||
class SMPL(nn.Module):
|
||||
"""SMPL 3d human mesh model of paper ref: Matthew Loper. ``SMPL: A skinned
|
||||
multi-person linear model''. This module is based on the smplx project
|
||||
(https://github.com/vchoutas/smplx).
|
||||
|
||||
Args:
|
||||
smpl_path (str): The path to the folder where the model weights are
|
||||
stored.
|
||||
joints_regressor (str): The path to the file where the joints
|
||||
regressor weight are stored.
|
||||
"""
|
||||
|
||||
def __init__(self, smpl_path, joints_regressor):
|
||||
super().__init__()
|
||||
|
||||
assert has_smpl, 'Please install smplx to use SMPL.'
|
||||
|
||||
self.smpl_neutral = SMPL_(
|
||||
model_path=smpl_path,
|
||||
create_global_orient=False,
|
||||
create_body_pose=False,
|
||||
create_transl=False,
|
||||
gender='neutral')
|
||||
|
||||
self.smpl_male = SMPL_(
|
||||
model_path=smpl_path,
|
||||
create_betas=False,
|
||||
create_global_orient=False,
|
||||
create_body_pose=False,
|
||||
create_transl=False,
|
||||
gender='male')
|
||||
|
||||
self.smpl_female = SMPL_(
|
||||
model_path=smpl_path,
|
||||
create_betas=False,
|
||||
create_global_orient=False,
|
||||
create_body_pose=False,
|
||||
create_transl=False,
|
||||
gender='female')
|
||||
|
||||
joints_regressor = torch.tensor(
|
||||
np.load(joints_regressor), dtype=torch.float)[None, ...]
|
||||
self.register_buffer('joints_regressor', joints_regressor)
|
||||
|
||||
self.num_verts = self.smpl_neutral.get_num_verts()
|
||||
self.num_joints = self.joints_regressor.shape[1]
|
||||
|
||||
def smpl_forward(self, model, **kwargs):
|
||||
"""Apply a specific SMPL model with given model parameters.
|
||||
|
||||
Note:
|
||||
B: batch size
|
||||
V: number of vertices
|
||||
K: number of joints
|
||||
|
||||
Returns:
|
||||
outputs (dict): Dict with mesh vertices and joints.
|
||||
- vertices: Tensor([B, V, 3]), mesh vertices
|
||||
- joints: Tensor([B, K, 3]), 3d joints regressed
|
||||
from mesh vertices.
|
||||
"""
|
||||
|
||||
betas = kwargs['betas']
|
||||
batch_size = betas.shape[0]
|
||||
device = betas.device
|
||||
output = {}
|
||||
if batch_size == 0:
|
||||
output['vertices'] = betas.new_zeros([0, self.num_verts, 3])
|
||||
output['joints'] = betas.new_zeros([0, self.num_joints, 3])
|
||||
else:
|
||||
smpl_out = model(**kwargs)
|
||||
output['vertices'] = smpl_out.vertices
|
||||
output['joints'] = torch.matmul(
|
||||
self.joints_regressor.to(device), output['vertices'])
|
||||
return output
|
||||
|
||||
def get_faces(self):
|
||||
"""Return mesh faces.
|
||||
|
||||
Note:
|
||||
F: number of faces
|
||||
|
||||
Returns:
|
||||
faces: np.ndarray([F, 3]), mesh faces
|
||||
"""
|
||||
return self.smpl_neutral.faces
|
||||
|
||||
def forward(self,
|
||||
betas,
|
||||
body_pose,
|
||||
global_orient,
|
||||
transl=None,
|
||||
gender=None):
|
||||
"""Forward function.
|
||||
|
||||
Note:
|
||||
B: batch size
|
||||
J: number of controllable joints of model, for smpl model J=23
|
||||
K: number of joints
|
||||
|
||||
Args:
|
||||
betas: Tensor([B, 10]), human body shape parameters of SMPL model.
|
||||
body_pose: Tensor([B, J*3] or [B, J, 3, 3]), human body pose
|
||||
parameters of SMPL model. It should be axis-angle vector
|
||||
([B, J*3]) or rotation matrix ([B, J, 3, 3)].
|
||||
global_orient: Tensor([B, 3] or [B, 1, 3, 3]), global orientation
|
||||
of human body. It should be axis-angle vector ([B, 3]) or
|
||||
rotation matrix ([B, 1, 3, 3)].
|
||||
transl: Tensor([B, 3]), global translation of human body.
|
||||
gender: Tensor([B]), gender parameters of human body. -1 for
|
||||
neutral, 0 for male , 1 for female.
|
||||
|
||||
Returns:
|
||||
outputs (dict): Dict with mesh vertices and joints.
|
||||
- vertices: Tensor([B, V, 3]), mesh vertices
|
||||
- joints: Tensor([B, K, 3]), 3d joints regressed from
|
||||
mesh vertices.
|
||||
"""
|
||||
|
||||
batch_size = betas.shape[0]
|
||||
pose2rot = True if body_pose.dim() == 2 else False
|
||||
if batch_size > 0 and gender is not None:
|
||||
output = {
|
||||
'vertices': betas.new_zeros([batch_size, self.num_verts, 3]),
|
||||
'joints': betas.new_zeros([batch_size, self.num_joints, 3])
|
||||
}
|
||||
|
||||
mask = gender < 0
|
||||
_out = self.smpl_forward(
|
||||
self.smpl_neutral,
|
||||
betas=betas[mask],
|
||||
body_pose=body_pose[mask],
|
||||
global_orient=global_orient[mask],
|
||||
transl=transl[mask] if transl is not None else None,
|
||||
pose2rot=pose2rot)
|
||||
output['vertices'][mask] = _out['vertices']
|
||||
output['joints'][mask] = _out['joints']
|
||||
|
||||
mask = gender == 0
|
||||
_out = self.smpl_forward(
|
||||
self.smpl_male,
|
||||
betas=betas[mask],
|
||||
body_pose=body_pose[mask],
|
||||
global_orient=global_orient[mask],
|
||||
transl=transl[mask] if transl is not None else None,
|
||||
pose2rot=pose2rot)
|
||||
output['vertices'][mask] = _out['vertices']
|
||||
output['joints'][mask] = _out['joints']
|
||||
|
||||
mask = gender == 1
|
||||
_out = self.smpl_forward(
|
||||
self.smpl_male,
|
||||
betas=betas[mask],
|
||||
body_pose=body_pose[mask],
|
||||
global_orient=global_orient[mask],
|
||||
transl=transl[mask] if transl is not None else None,
|
||||
pose2rot=pose2rot)
|
||||
output['vertices'][mask] = _out['vertices']
|
||||
output['joints'][mask] = _out['joints']
|
||||
else:
|
||||
return self.smpl_forward(
|
||||
self.smpl_neutral,
|
||||
betas=betas,
|
||||
body_pose=body_pose,
|
||||
global_orient=global_orient,
|
||||
transl=transl,
|
||||
pose2rot=pose2rot)
|
||||
|
||||
return output
|
||||
Reference in New Issue
Block a user