Spaces:
Running
Running
| import os | |
| import numpy as np | |
| import torch | |
| from torch.nn import Module | |
| from custom_manopth.smpl_handpca_wrapper_HAND_only import ready_arguments | |
| from custom_manopth import rodrigues_layer, rotproj, rot6d | |
| from custom_manopth.tensutils import (th_posemap_axisang, th_with_zeros, th_pack, | |
| subtract_flat_id, make_list) | |
| class ManoLayer(Module): | |
| __constants__ = [ | |
| 'use_pca', 'rot', 'ncomps', 'ncomps', 'kintree_parents', 'check', | |
| 'side', 'center_idx', 'joint_rot_mode' | |
| ] | |
| def __init__(self, | |
| center_idx=None, | |
| flat_hand_mean=True, | |
| ncomps=6, | |
| side='right', | |
| mano_root='mano/models', | |
| use_pca=True, | |
| root_rot_mode='axisang', | |
| joint_rot_mode='axisang', | |
| robust_rot=False): | |
| """ | |
| Args: | |
| center_idx: index of center joint in our computations, | |
| if -1 centers on estimate of palm as middle of base | |
| of middle finger and wrist | |
| flat_hand_mean: if True, (0, 0, 0, ...) pose coefficients match | |
| flat hand, else match average hand pose | |
| mano_root: path to MANO pkl files for left and right hand | |
| ncomps: number of PCA components form pose space (<45) | |
| side: 'right' or 'left' | |
| use_pca: Use PCA decomposition for pose space. | |
| joint_rot_mode: 'axisang' or 'rotmat', ignored if use_pca | |
| """ | |
| super().__init__() | |
| self.center_idx = center_idx | |
| self.robust_rot = robust_rot | |
| if root_rot_mode == 'axisang': | |
| self.rot = 3 | |
| else: | |
| self.rot = 6 | |
| self.flat_hand_mean = flat_hand_mean | |
| self.side = side | |
| self.use_pca = use_pca | |
| self.joint_rot_mode = joint_rot_mode | |
| self.root_rot_mode = root_rot_mode | |
| if use_pca: | |
| self.ncomps = ncomps | |
| else: | |
| self.ncomps = 45 | |
| if side == 'right': | |
| self.mano_path = os.path.join(mano_root, 'MANO_RIGHT.pkl') | |
| elif side == 'left': | |
| self.mano_path = os.path.join(mano_root, 'MANO_LEFT.pkl') | |
| smpl_data = ready_arguments(self.mano_path) | |
| hands_components = smpl_data['hands_components'] | |
| self.smpl_data = smpl_data | |
| self.register_buffer('th_betas', | |
| torch.Tensor(smpl_data['betas']).unsqueeze(0)) | |
| self.register_buffer('th_shapedirs', | |
| torch.Tensor(smpl_data['shapedirs'])) | |
| self.register_buffer('th_posedirs', | |
| torch.Tensor(smpl_data['posedirs'])) | |
| self.register_buffer( | |
| 'th_v_template', | |
| torch.Tensor(smpl_data['v_template']).unsqueeze(0)) | |
| self.register_buffer( | |
| 'th_J_regressor', | |
| torch.Tensor(np.array(smpl_data['J_regressor'].toarray()))) | |
| self.register_buffer('th_weights', | |
| torch.Tensor(smpl_data['weights'])) | |
| self.register_buffer('th_faces', | |
| torch.Tensor(smpl_data['f'].astype(np.int32)).long()) | |
| # Get hand mean | |
| hands_mean = np.zeros(hands_components.shape[1] | |
| ) if flat_hand_mean else smpl_data['hands_mean'] | |
| hands_mean = hands_mean.copy() | |
| th_hands_mean = torch.Tensor(hands_mean).unsqueeze(0) | |
| if self.use_pca or self.joint_rot_mode == 'axisang': | |
| # Save as axis-angle | |
| self.register_buffer('th_hands_mean', th_hands_mean) | |
| selected_components = hands_components[:ncomps] | |
| self.register_buffer('th_comps', torch.Tensor(hands_components)) | |
| self.register_buffer('th_selected_comps', | |
| torch.Tensor(selected_components)) | |
| else: | |
| th_hands_mean_rotmat = rodrigues_layer.batch_rodrigues( | |
| th_hands_mean.view(15, 3)).reshape(15, 3, 3) | |
| self.register_buffer('th_hands_mean_rotmat', th_hands_mean_rotmat) | |
| # Kinematic chain params | |
| self.kintree_table = smpl_data['kintree_table'] | |
| parents = list(self.kintree_table[0].tolist()) | |
| self.kintree_parents = parents | |
| def forward(self, | |
| th_pose_coeffs, | |
| th_betas=torch.zeros(1), | |
| th_trans=torch.zeros(1), | |
| root_palm=torch.Tensor([0]), | |
| share_betas=torch.Tensor([0]), | |
| ): | |
| """ | |
| Args: | |
| th_trans (Tensor (batch_size x ncomps)): if provided, applies trans to joints and vertices | |
| th_betas (Tensor (batch_size x 10)): if provided, uses given shape parameters for hand shape | |
| else centers on root joint (9th joint) | |
| root_palm: return palm as hand root instead of wrist | |
| """ | |
| # if len(th_pose_coeffs) == 0: | |
| # return th_pose_coeffs.new_empty(0), th_pose_coeffs.new_empty(0) | |
| batch_size = th_pose_coeffs.shape[0] | |
| # Get axis angle from PCA components and coefficients | |
| if self.use_pca or self.joint_rot_mode == 'axisang': | |
| # Remove global rot coeffs | |
| th_hand_pose_coeffs = th_pose_coeffs[:, self.rot:self.rot + | |
| self.ncomps] | |
| if self.use_pca: | |
| # PCA components --> axis angles | |
| th_full_hand_pose = th_hand_pose_coeffs.mm(self.th_selected_comps) | |
| else: | |
| th_full_hand_pose = th_hand_pose_coeffs | |
| # Concatenate back global rot | |
| th_full_pose = torch.cat([ | |
| th_pose_coeffs[:, :self.rot], | |
| self.th_hands_mean + th_full_hand_pose | |
| ], 1) | |
| if self.root_rot_mode == 'axisang': | |
| # compute rotation matrixes from axis-angle while skipping global rotation | |
| th_pose_map, th_rot_map = th_posemap_axisang(th_full_pose) | |
| root_rot = th_rot_map[:, :9].view(batch_size, 3, 3) | |
| th_rot_map = th_rot_map[:, 9:] | |
| th_pose_map = th_pose_map[:, 9:] | |
| else: | |
| # th_posemap offsets by 3, so add offset or 3 to get to self.rot=6 | |
| th_pose_map, th_rot_map = th_posemap_axisang(th_full_pose[:, 6:]) | |
| if self.robust_rot: | |
| root_rot = rot6d.robust_compute_rotation_matrix_from_ortho6d(th_full_pose[:, :6]) | |
| else: | |
| root_rot = rot6d.compute_rotation_matrix_from_ortho6d(th_full_pose[:, :6]) | |
| else: | |
| assert th_pose_coeffs.dim() == 4, ( | |
| 'When not self.use_pca, ' | |
| 'th_pose_coeffs should have 4 dims, got {}'.format( | |
| th_pose_coeffs.dim())) | |
| assert th_pose_coeffs.shape[2:4] == (3, 3), ( | |
| 'When not self.use_pca, th_pose_coeffs have 3x3 matrix for two' | |
| 'last dims, got {}'.format(th_pose_coeffs.shape[2:4])) | |
| th_pose_rots = rotproj.batch_rotprojs(th_pose_coeffs) | |
| th_rot_map = th_pose_rots[:, 1:].view(batch_size, -1) | |
| th_pose_map = subtract_flat_id(th_rot_map) | |
| root_rot = th_pose_rots[:, 0] | |
| # Full axis angle representation with root joint | |
| if th_betas is None or th_betas.numel() == 1: | |
| th_v_shaped = torch.matmul(self.th_shapedirs, | |
| self.th_betas.transpose(1, 0)).permute( | |
| 2, 0, 1) + self.th_v_template | |
| th_j = torch.matmul(self.th_J_regressor, th_v_shaped).repeat( | |
| batch_size, 1, 1) | |
| else: | |
| if share_betas: | |
| th_betas = th_betas.mean(0, keepdim=True).expand(th_betas.shape[0], 10) | |
| th_v_shaped = torch.matmul(self.th_shapedirs, | |
| th_betas.transpose(1, 0)).permute( | |
| 2, 0, 1) + self.th_v_template | |
| th_j = torch.matmul(self.th_J_regressor, th_v_shaped) | |
| # th_pose_map should have shape 20x135 | |
| th_v_posed = th_v_shaped + torch.matmul( | |
| self.th_posedirs, th_pose_map.transpose(0, 1)).permute(2, 0, 1) | |
| # Final T pose with transformation done ! | |
| # Global rigid transformation | |
| root_j = th_j[:, 0, :].contiguous().view(batch_size, 3, 1) | |
| root_trans = th_with_zeros(torch.cat([root_rot, root_j], 2)) | |
| all_rots = th_rot_map.view(th_rot_map.shape[0], 15, 3, 3) | |
| lev1_idxs = [1, 4, 7, 10, 13] | |
| lev2_idxs = [2, 5, 8, 11, 14] | |
| lev3_idxs = [3, 6, 9, 12, 15] | |
| lev1_rots = all_rots[:, [idx - 1 for idx in lev1_idxs]] | |
| lev2_rots = all_rots[:, [idx - 1 for idx in lev2_idxs]] | |
| lev3_rots = all_rots[:, [idx - 1 for idx in lev3_idxs]] | |
| lev1_j = th_j[:, lev1_idxs] | |
| lev2_j = th_j[:, lev2_idxs] | |
| lev3_j = th_j[:, lev3_idxs] | |
| # From base to tips | |
| # Get lev1 results | |
| all_transforms = [root_trans.unsqueeze(1)] | |
| lev1_j_rel = lev1_j - root_j.transpose(1, 2) | |
| lev1_rel_transform_flt = th_with_zeros(torch.cat([lev1_rots, lev1_j_rel.unsqueeze(3)], 3).view(-1, 3, 4)) | |
| root_trans_flt = root_trans.unsqueeze(1).repeat(1, 5, 1, 1).view(root_trans.shape[0] * 5, 4, 4) | |
| lev1_flt = torch.matmul(root_trans_flt, lev1_rel_transform_flt) | |
| all_transforms.append(lev1_flt.view(all_rots.shape[0], 5, 4, 4)) | |
| # Get lev2 results | |
| lev2_j_rel = lev2_j - lev1_j | |
| lev2_rel_transform_flt = th_with_zeros(torch.cat([lev2_rots, lev2_j_rel.unsqueeze(3)], 3).view(-1, 3, 4)) | |
| lev2_flt = torch.matmul(lev1_flt, lev2_rel_transform_flt) | |
| all_transforms.append(lev2_flt.view(all_rots.shape[0], 5, 4, 4)) | |
| # Get lev3 results | |
| lev3_j_rel = lev3_j - lev2_j | |
| lev3_rel_transform_flt = th_with_zeros(torch.cat([lev3_rots, lev3_j_rel.unsqueeze(3)], 3).view(-1, 3, 4)) | |
| lev3_flt = torch.matmul(lev2_flt, lev3_rel_transform_flt) | |
| all_transforms.append(lev3_flt.view(all_rots.shape[0], 5, 4, 4)) | |
| reorder_idxs = [0, 1, 6, 11, 2, 7, 12, 3, 8, 13, 4, 9, 14, 5, 10, 15] | |
| th_results = torch.cat(all_transforms, 1)[:, reorder_idxs] | |
| th_results_global = th_results | |
| joint_js = torch.cat([th_j, th_j.new_zeros(th_j.shape[0], 16, 1)], 2) | |
| tmp2 = torch.matmul(th_results, joint_js.unsqueeze(3)) | |
| th_results2 = (th_results - torch.cat([tmp2.new_zeros(*tmp2.shape[:2], 4, 3), tmp2], 3)).permute(0, 2, 3, 1) | |
| th_T = torch.matmul(th_results2, self.th_weights.transpose(0, 1)) | |
| th_rest_shape_h = torch.cat([ | |
| th_v_posed.transpose(2, 1), | |
| torch.ones((batch_size, 1, th_v_posed.shape[1]), | |
| dtype=th_T.dtype, | |
| device=th_T.device), | |
| ], 1) | |
| th_verts = (th_T * th_rest_shape_h.unsqueeze(1)).sum(2).transpose(2, 1) | |
| th_verts = th_verts[:, :, :3] | |
| th_jtr = th_results_global[:, :, :3, 3] | |
| # In addition to MANO reference joints we sample vertices on each finger | |
| # to serve as finger tips | |
| if self.side == 'right': | |
| tips = th_verts[:, [745, 317, 444, 556, 673]] | |
| else: | |
| tips = th_verts[:, [745, 317, 445, 556, 673]] | |
| if bool(root_palm): | |
| palm = (th_verts[:, 95] + th_verts[:, 22]).unsqueeze(1) / 2 | |
| th_jtr = torch.cat([palm, th_jtr[:, 1:]], 1) | |
| th_jtr = torch.cat([th_jtr, tips], 1) | |
| # Reorder joints to match visualization utilities | |
| th_jtr = th_jtr[:, [0, 13, 14, 15, 16, 1, 2, 3, 17, 4, 5, 6, 18, 10, 11, 12, 19, 7, 8, 9, 20]] | |
| if th_trans is None or bool(torch.norm(th_trans) == 0): | |
| if self.center_idx is not None: | |
| center_joint = th_jtr[:, self.center_idx].unsqueeze(1) | |
| th_jtr = th_jtr - center_joint | |
| th_verts = th_verts - center_joint | |
| else: | |
| th_jtr = th_jtr + th_trans.unsqueeze(1) | |
| th_verts = th_verts + th_trans.unsqueeze(1) | |
| # Scale to milimeters | |
| th_verts = th_verts * 1000 | |
| th_jtr = th_jtr * 1000 | |
| return th_verts, th_jtr | |