|
|
import torch |
|
|
import hydra |
|
|
import numpy as np |
|
|
from ..smpl.body_models import SMPL |
|
|
|
|
|
class SMPLServer(torch.nn.Module): |
|
|
|
|
|
def __init__(self, gender='neutral', betas=None, v_template=None): |
|
|
super().__init__() |
|
|
|
|
|
|
|
|
self.smpl = SMPL(model_path=hydra.utils.to_absolute_path('lib/smpl/smpl_model'), |
|
|
gender=gender, |
|
|
batch_size=1, |
|
|
use_hands=False, |
|
|
use_feet_keypoints=False, |
|
|
dtype=torch.float32).cuda() |
|
|
|
|
|
self.bone_parents = self.smpl.bone_parents.astype(int) |
|
|
self.bone_parents[0] = -1 |
|
|
self.bone_ids = [] |
|
|
self.faces = self.smpl.faces |
|
|
for i in range(24): self.bone_ids.append([self.bone_parents[i], i]) |
|
|
|
|
|
if v_template is not None: |
|
|
self.v_template = torch.tensor(v_template).float().cuda() |
|
|
else: |
|
|
self.v_template = None |
|
|
|
|
|
if betas is not None: |
|
|
self.betas = torch.tensor(betas).float().cuda() |
|
|
else: |
|
|
self.betas = None |
|
|
|
|
|
|
|
|
param_canonical = torch.zeros((1, 86),dtype=torch.float32).cuda() |
|
|
param_canonical[0, 0] = 1 |
|
|
param_canonical[0, 9] = np.pi / 6 |
|
|
param_canonical[0, 12] = -np.pi / 6 |
|
|
if self.betas is not None and self.v_template is None: |
|
|
param_canonical[0,-10:] = self.betas |
|
|
self.param_canonical = param_canonical |
|
|
|
|
|
output = self.forward(*torch.split(self.param_canonical, [1, 3, 72, 10], dim=1), absolute=True) |
|
|
self.verts_c = output['smpl_verts'] |
|
|
self.joints_c = output['smpl_jnts'] |
|
|
self.tfs_c_inv = output['smpl_tfs'].squeeze(0).inverse() |
|
|
|
|
|
|
|
|
def forward(self, scale, transl, thetas, betas, absolute=False): |
|
|
"""return SMPL output from params |
|
|
Args: |
|
|
scale : scale factor. shape: [B, 1] |
|
|
transl: translation. shape: [B, 3] |
|
|
thetas: pose. shape: [B, 72] |
|
|
betas: shape. shape: [B, 10] |
|
|
absolute (bool): if true return smpl_tfs wrt thetas=0. else wrt thetas=thetas_canonical. |
|
|
Returns: |
|
|
smpl_verts: vertices. shape: [B, 6893. 3] |
|
|
smpl_tfs: bone transformations. shape: [B, 24, 4, 4] |
|
|
smpl_jnts: joint positions. shape: [B, 25, 3] |
|
|
""" |
|
|
|
|
|
output = {} |
|
|
|
|
|
|
|
|
if self.v_template is not None: |
|
|
betas = torch.zeros_like(betas) |
|
|
|
|
|
|
|
|
smpl_output = self.smpl.forward(betas=betas, |
|
|
transl=torch.zeros_like(transl), |
|
|
body_pose=thetas[:, 3:], |
|
|
global_orient=thetas[:, :3], |
|
|
return_verts=True, |
|
|
return_full_pose=True, |
|
|
v_template=self.v_template) |
|
|
|
|
|
verts = smpl_output.vertices.clone() |
|
|
output['smpl_verts'] = verts * scale.unsqueeze(1) + transl.unsqueeze(1) * scale.unsqueeze(1) |
|
|
|
|
|
joints = smpl_output.joints.clone() |
|
|
output['smpl_jnts'] = joints * scale.unsqueeze(1) + transl.unsqueeze(1) * scale.unsqueeze(1) |
|
|
|
|
|
tf_mats = smpl_output.T.clone() |
|
|
tf_mats[:, :, :3, :] = tf_mats[:, :, :3, :] * scale.unsqueeze(1).unsqueeze(1) |
|
|
tf_mats[:, :, :3, 3] = tf_mats[:, :, :3, 3] + transl.unsqueeze(1) * scale.unsqueeze(1) |
|
|
|
|
|
if not absolute: |
|
|
tf_mats = torch.einsum('bnij,njk->bnik', tf_mats, self.tfs_c_inv) |
|
|
|
|
|
output['smpl_tfs'] = tf_mats |
|
|
output['smpl_weights'] = smpl_output.weights |
|
|
return output |
