""" Copyright©2024 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute for Intelligent Systems. All rights reserved. Author: Marilyn Keller See https://skel.is.tue.mpg.de/license.html for licensing and contact information. """ import os import torch.nn as nn import torch import numpy as np import pickle as pkl from typing import NewType, Optional from lib.body_models.skel.joints_def import curve_torch_3d, left_scapula, right_scapula from lib.body_models.skel.osim_rot import ConstantCurvatureJoint, CustomJoint, EllipsoidJoint, PinJoint, WalkerKnee from lib.body_models.skel.utils import build_homog_matrix, rotation_matrix_from_vectors, sparce_coo_matrix2tensor, with_zeros, matmul_chain from dataclasses import dataclass, fields from lib.body_models.skel.kin_skel import scaling_keypoints, pose_param_names, smpl_joint_corresp import lib.body_models.skel.config as cg Tensor = NewType('Tensor', torch.Tensor) @dataclass class ModelOutput: vertices: Optional[Tensor] = None joints: Optional[Tensor] = None full_pose: Optional[Tensor] = None global_orient: Optional[Tensor] = None transl: Optional[Tensor] = None v_shaped: Optional[Tensor] = None def __getitem__(self, key): return getattr(self, key) def get(self, key, default=None): return getattr(self, key, default) def __iter__(self): return self.keys() def keys(self): keys = [t.name for t in fields(self)] return iter(keys) def values(self): values = [getattr(self, t.name) for t in fields(self)] return iter(values) def items(self): data = [(t.name, getattr(self, t.name)) for t in fields(self)] return iter(data) @dataclass class SKELOutput(ModelOutput): betas: Optional[Tensor] = None body_pose: Optional[Tensor] = None skin_verts: Optional[Tensor] = None skel_verts: Optional[Tensor] = None joints: Optional[Tensor] = None joints_ori: Optional[Tensor] = None betas: Optional[Tensor] = None poses: Optional[Tensor] = None trans : Optional[Tensor] = None pose_offsets : Optional[Tensor] = None joints_tpose : Optional[Tensor] = None v_skin_shaped : Optional[Tensor] = None class SKEL(nn.Module): num_betas = 10 def __init__(self, gender, model_path=None, custom_joint_reg_path=None, **kwargs): super(SKEL, self).__init__() if gender not in ['male', 'female']: raise RuntimeError(f'Invalid Gender, got {gender}') self.gender = gender if model_path is None: # skel_file = f"/Users/mkeller2/Data/skel_models_v1.0/skel_{gender}.pkl" skel_file = os.path.join(cg.skel_folder, f"skel_{gender}.pkl") else: skel_file = os.path.join(model_path, f"skel_{gender}.pkl") assert os.path.exists(skel_file), f"Skel model file {skel_file} does not exist" skel_data = pkl.load(open(skel_file, 'rb')) # Check that the version of the skel model is compatible with this loader assert 'version' in skel_data, f"Expected version 1.1.1 of the SKEL picke. Please download the latest skel pkl versions from https://skel.is.tue.mpg.de/download.html" version = skel_data['version'] assert version == '1.1.1', f"Expected version 1.1.1, got {version}. Please download the latest skel pkl versions from https://skel.is.tue.mpg.de/download.html" self.num_betas = 10 self.num_q_params = 46 self.bone_names = skel_data['bone_names'] self.num_joints = skel_data['J_regressor_osim'].shape[0] self.num_joints_smpl = skel_data['J_regressor'].shape[0] self.joints_name = skel_data['joints_name'] self.pose_params_name = skel_data['pose_params_name'] # register the template meshes self.register_buffer('skin_template_v', torch.FloatTensor(skel_data['skin_template_v'])) self.register_buffer('skin_f', torch.LongTensor(skel_data['skin_template_f'])) self.register_buffer('skel_template_v', torch.FloatTensor(skel_data['skel_template_v'])) self.register_buffer('skel_f', torch.LongTensor(skel_data['skel_template_f'])) # Shape corrective blend shapes self.register_buffer('shapedirs', torch.FloatTensor(np.array(skel_data['shapedirs'][:,:,:self.num_betas]))) self.register_buffer('posedirs', torch.FloatTensor(np.array(skel_data['posedirs']))) # Model sparse joints regressor, regresses joints location from a mesh self.register_buffer('J_regressor', sparce_coo_matrix2tensor(skel_data['J_regressor'])) # Regress the anatomical joint location with a regressor learned from BioAmass if custom_joint_reg_path is not None: J_regressor_skel = pkl.load(open(custom_joint_reg_path, 'rb')) if 'scipy.sparse' in str(type(J_regressor_skel)): J_regressor_skel = J_regressor_skel.todense() self.register_buffer('J_regressor_osim', torch.FloatTensor(J_regressor_skel)) print('WARNING: Using custom joint regressor') else: self.register_buffer('J_regressor_osim', sparce_coo_matrix2tensor(skel_data['J_regressor_osim'], make_dense=True)) self.register_buffer('per_joint_rot', torch.FloatTensor(skel_data['per_joint_rot'])) # Skin model skinning weights self.register_buffer('skin_weights', sparce_coo_matrix2tensor(skel_data['skin_weights'])) # Skeleton model skinning weights self.register_buffer('skel_weights', sparce_coo_matrix2tensor(skel_data['skel_weights'])) self.register_buffer('skel_weights_rigid', sparce_coo_matrix2tensor(skel_data['skel_weights_rigid'])) # Kinematic tree of the model self.register_buffer('kintree_table', torch.from_numpy(skel_data['osim_kintree_table'].astype(np.int64))) self.register_buffer('parameter_mapping', torch.from_numpy(skel_data['parameter_mapping'].astype(np.int64))) # transformation from osim can pose to T pose self.register_buffer('tpose_transfo', torch.FloatTensor(skel_data['tpose_transfo'])) # transformation from osim can pose to A pose self.register_buffer('apose_transfo', torch.FloatTensor(skel_data['apose_transfo'])) self.register_buffer('apose_rel_transfo', torch.FloatTensor(skel_data['apose_rel_transfo'])) # Indices of bones which orientation should not vary with beta in T pose: joint_idx_fixed_beta = [0, 5, 10, 13, 18, 23] self.register_buffer('joint_idx_fixed_beta', torch.IntTensor(joint_idx_fixed_beta)) id_to_col = {self.kintree_table[1, i].item(): i for i in range(self.kintree_table.shape[1])} self.register_buffer('parent', torch.LongTensor( [id_to_col[self.kintree_table[0, it].item()] for it in range(1, self.kintree_table.shape[1])])) # child array # TODO create this array in the SKEL creator child_array = [] Nj = self.num_joints for i in range(0, Nj): try: j_array = torch.where(self.kintree_table[0] == i)[0] # candidate child lines if len(j_array) == 0: child_index = 0 else: j = j_array[0] if j>=len(self.kintree_table[1]): child_index = 0 else: child_index = self.kintree_table[1,j].item() child_array.append(child_index) except: import ipdb; ipdb.set_trace() # print(f"child_array: ") # [print(i,child_array[i]) for i in range(0, Nj)] self.register_buffer('child', torch.LongTensor(child_array)) # Instantiate joints self.joints_dict = nn.ModuleList([ CustomJoint(axis=[[0,0,1], [1,0,0], [0,1,0]], axis_flip=[1, 1, 1]), # 0 pelvis CustomJoint(axis=[[0,0,1], [1,0,0], [0,1,0]], axis_flip=[1, 1, 1]), # 1 femur_r WalkerKnee(), # 2 tibia_r PinJoint(parent_frame_ori = [0.175895, -0.105208, 0.0186622]), # 3 talus_r Field taken from .osim Joint-> frames -> PhysicalOffsetFrame -> orientation PinJoint(parent_frame_ori = [-1.76818999, 0.906223, 1.8196000]), # 4 calcn_r PinJoint(parent_frame_ori = [-3.141589999, 0.6199010, 0]), # 5 toes_r CustomJoint(axis=[[0,0,1], [1,0,0], [0,1,0]], axis_flip=[1, -1, -1]), # 6 femur_l WalkerKnee(), # 7 tibia_l PinJoint(parent_frame_ori = [0.175895, -0.105208, 0.0186622]), # 8 talus_l PinJoint(parent_frame_ori = [1.768189999 ,-0.906223, 1.8196000]), # 9 calcn_l PinJoint(parent_frame_ori = [-3.141589999, -0.6199010, 0]), # 10 toes_l ConstantCurvatureJoint(axis=[[1,0,0], [0,0,1], [0,1,0]], axis_flip=[1, 1, 1]), # 11 lumbar ConstantCurvatureJoint(axis=[[1,0,0], [0,0,1], [0,1,0]], axis_flip=[1, 1, 1]), # 12 thorax ConstantCurvatureJoint(axis=[[1,0,0], [0,0,1], [0,1,0]], axis_flip=[1, 1, 1]), # 13 head EllipsoidJoint(axis=[[0,1,0], [0,0,1], [1,0,0]], axis_flip=[1, -1, -1]), # 14 scapula_r CustomJoint(axis=[[1,0,0], [0,1,0], [0,0,1]], axis_flip=[1, 1, 1]), # 15 humerus_r CustomJoint(axis=[[0.0494, 0.0366, 0.99810825]], axis_flip=[[1]]), # 16 ulna_r CustomJoint(axis=[[-0.01716099, 0.99266564, -0.11966796]], axis_flip=[[1]]), # 17 radius_r CustomJoint(axis=[[1,0,0], [0,0,-1]], axis_flip=[1, 1]), # 18 hand_r EllipsoidJoint(axis=[[0,1,0], [0,0,1], [1,0,0]], axis_flip=[1, 1, 1]), # 19 scapula_l CustomJoint(axis=[[1,0,0], [0,1,0], [0,0,1]], axis_flip=[1, 1, 1]), # 20 humerus_l CustomJoint(axis=[[-0.0494, -0.0366, 0.99810825]], axis_flip=[[1]]), # 21 ulna_l CustomJoint(axis=[[0.01716099, -0.99266564, -0.11966796]], axis_flip=[[1]]), # 22 radius_l CustomJoint(axis=[[-1,0,0], [0,0,-1]], axis_flip=[1, 1]), # 23 hand_l ]) def pose_params_to_rot(self, osim_poses): """ Transform the pose parameters to 3x3 rotation matrices Each parameter is mapped to a joint as described in joint_dict. The specific joint object is then used to compute the rotation matrix. """ B = osim_poses.shape[0] Nj = self.num_joints ident = torch.eye(3, dtype=osim_poses.dtype).to(osim_poses.device) Rp = ident.unsqueeze(0).unsqueeze(0).repeat(B, Nj,1,1) tp = torch.zeros(B, Nj, 3).to(osim_poses.device) start_index = 0 for i in range(0, Nj): joint_object = self.joints_dict[i] end_index = start_index + joint_object.nb_dof Rp[:, i] = joint_object.q_to_rot(osim_poses[:, start_index:end_index]) start_index = end_index return Rp, tp def params_name_to_index(self, param_name): assert param_name in pose_param_names param_index = pose_param_names.index(param_name) return param_index def forward(self, poses, betas, trans, poses_type='skel', skelmesh=True, dJ=None, pose_dep_bs=True): """ params poses : B x 46 tensor of pose parameters betas : B x 10 tensor of shape parameters, same as SMPL trans : B x 3 tensor of translation poses_type : str, 'skel', should not be changed skelmesh : bool, if True, returns the skeleton vertices. The skeleton mesh is heavy so to fit on GPU memory, set to False when not needed. dJ : B x 24 x 3 tensor of the offset of the joints location from the anatomical regressor. If None, the offset is set to 0. pose_dep_bs : bool, if True (default), applies the pose dependant blend shapes. If False, the pose dependant blend shapes are not applied. return SKELOutput class with the following fields: betas : Bx10 tensor of shape parameters poses : Bx46 tensor of pose parameters skin_verts : Bx6890x3 tensor of skin vertices skel_verts : tensor of skeleton vertices joints : Bx24x3 tensor of joints location joints_ori : Bx24x3x3 tensor of joints orientation trans : Bx3 pose dependant blend shapes offsets pose_offsets : Bx6080x3 pose dependant blend shapes offsets joints_tpose : Bx24x3 3D joints location in T pose In this function we use the following conventions: B : batch size Ns : skin vertices Nk : skeleton vertices """ Ns = self.skin_template_v.shape[0] # nb skin vertices Nk = self.skel_template_v.shape[0] # nb skeleton vertices Nj = self.num_joints B = poses.shape[0] device = poses.device # Check the shapes of the inputs assert len(betas.shape) == 2, f"Betas should be of shape (B, {self.num_betas}), but got {betas.shape}" assert poses.shape[0] == betas.shape[0], f"Expected poses and betas to have the same batch size, but got {poses.shape[0]} and {betas.shape[0]}" assert poses.shape[0] == trans.shape[0], f"Expected poses and betas to have the same batch size, but got {poses.shape[0]} and {trans.shape[0]}" if dJ is not None: assert len(dJ.shape) == 3, f"Expected dJ to have shape (B, {Nj}, 3), but got {dJ.shape}" assert dJ is None or dJ.shape[0] == B, f"Expected dJ to have the same batch size as poses, but got {dJ.shape[0]} and {poses.shape[0]}" assert dJ.shape[1] == Nj, f"Expected dJ to have the same number of joints as the model, but got {dJ.shape[1]} and {Nj}" # Check the device of the inputs assert betas.device == device, f"Betas should be on device {device}, but got {betas.device}" assert trans.device == device, f"Trans should be on device {device}, but got {trans.device}" skin_v0 = self.skin_template_v[None, :] skel_v0 = self.skel_template_v[None, :] betas = betas[:, :, None] # TODO Name the expanded beta differently # TODO clean this part assert poses_type in ['skel', 'bsm'], f"got {poses_type}" if poses_type == 'bsm': assert poses.shape[1] == self.num_q_params - 3, f'With poses_type bsm, expected parameters of shape (B, {self.num_q_params - 3}, got {poses.shape}' poses_bsm = poses poses_skel = torch.zeros(B, self.num_q_params) poses_skel[:,:3] = poses_bsm[:, :3] trans = poses_bsm[:, 3:6] # In BSM parametrization, the hips translation is given by params 3 to 5 poses_skel[:, 3:] = poses_bsm poses = poses_skel else: assert poses.shape[1] == self.num_q_params, f'With poses_type skel, expected parameters of shape (B, {self.num_q_params}), got {poses.shape}' pass # Load poses as expected # Distinction bsm skel. by default it will be bsm # ------- Shape ---------- # Apply the beta offset to the template shapedirs = self.shapedirs.view(-1, self.num_betas)[None, :].expand(B, -1, -1) # B x D*Ns x num_betas v_shaped = skin_v0 + torch.matmul(shapedirs, betas).view(B, Ns, 3) # ------- Joints ---------- # Regress the anatomical joint location J = torch.einsum('bik,ji->bjk', [v_shaped, self.J_regressor_osim]) # BxJx3 # osim regressor # J = self.apose_transfo[:, :3, -1].view(1, Nj, 3).expand(B, -1, -1) # Osim default pose joints location if dJ is not None: J = J + dJ J_tpose = J.clone() # Local translation J_ = J.clone() # BxJx3 J_[:, 1:, :] = J[:, 1:, :] - J[:, self.parent, :] t = J_[:, :, :, None] # BxJx3x1 # ------- Bones transformation matrix---------- # Bone initial transform to go from unposed to SMPL T pose Rk01 = self.compute_bone_orientation(J, J_) # BSM default pose rotations Ra = self.apose_rel_transfo[:, :3, :3].view(1, Nj, 3,3).expand(B, Nj, 3, 3) # Local bone rotation given by the pose param Rp, tp = self.pose_params_to_rot(poses) # BxNjx3x3 pose params to rotation R = matmul_chain([Rk01, Ra.transpose(2,3), Rp, Ra, Rk01.transpose(2,3)]) ###### Compute translation for non pure rotation joints t_posed = t.clone() # Scapula thorax_width = torch.norm(J[:, 19, :] - J[:, 14, :], dim=1) # Distance between the two scapula joints, size B thorax_height = torch.norm(J[:, 12, :] - J[:, 11, :], dim=1) # Distance between the two scapula joints, size B angle_abduction = poses[:,26] angle_elevation = poses[:,27] angle_rot = poses[:,28] angle_zero = torch.zeros_like(angle_abduction) t_posed[:,14] = t_posed[:,14] + \ (right_scapula(angle_abduction, angle_elevation, angle_rot, thorax_width, thorax_height).view(-1,3,1) - right_scapula(angle_zero, angle_zero, angle_zero, thorax_width, thorax_height).view(-1,3,1)) angle_abduction = poses[:,36] angle_elevation = poses[:,37] angle_rot = poses[:,38] angle_zero = torch.zeros_like(angle_abduction) t_posed[:,19] = t_posed[:,19] + \ (left_scapula(angle_abduction, angle_elevation, angle_rot, thorax_width, thorax_height).view(-1,3,1) - left_scapula(angle_zero, angle_zero, angle_zero, thorax_width, thorax_height).view(-1,3,1)) # Knee_r # TODO add the Walker knee offset # bone_scale = self.compute_bone_scale(J_,J, skin_v0, v_shaped) # f1 = poses[:, 2*3+2].clone() # scale_femur = bone_scale[:, 2] # factor = 0.076/0.080 * scale_femur # The template femur medial laterak spacing #66 # f = -f1*180/torch.pi #knee_flexion # varus = (0.12367*f)-0.0009*f**2 # introt = 0.3781*f-0.001781*f**2 # ydis = (-0.0683*f # + 8.804e-4 * f**2 # - 3.750e-06*f**3 # )/1000*factor # up-down # zdis = (-0.1283*f # + 4.796e-4 * f**2)/1000*factor # # import ipdb; ipdb.set_trace() # poses[:, 9] = poses[:, 9] + varus # t_posed[:,2] = t_posed[:,2] + torch.stack([torch.zeros_like(ydis), ydis, zdis], dim=1).view(-1,3,1) # poses[:, 2*3+2]=0 # t_unposed = torch.zeros_like(t_posed) # t_unposed[:,2] = torch.stack([torch.zeros_like(ydis), ydis, zdis], dim=1).view(-1,3,1) # Spine lumbar_bending = poses[:,17] lumbar_extension = poses[:,18] angle_zero = torch.zeros_like(lumbar_bending) interp_t = torch.ones_like(lumbar_bending) l = torch.abs(J[:, 11, 1] - J[:, 0, 1]) # Length of the spine section along y axis t_posed[:,11] = t_posed[:,11] + \ (curve_torch_3d(lumbar_bending, lumbar_extension, t=interp_t, l=l) - curve_torch_3d(angle_zero, angle_zero, t=interp_t, l=l)) thorax_bending = poses[:,20] thorax_extension = poses[:,21] angle_zero = torch.zeros_like(thorax_bending) interp_t = torch.ones_like(thorax_bending) l = torch.abs(J[:, 12, 1] - J[:, 11, 1]) # Length of the spine section t_posed[:,12] = t_posed[:,12] + \ (curve_torch_3d(thorax_bending, thorax_extension, t=interp_t, l=l) - curve_torch_3d(angle_zero, angle_zero, t=interp_t, l=l)) head_bending = poses[:, 23] head_extension = poses[:,24] angle_zero = torch.zeros_like(head_bending) interp_t = torch.ones_like(head_bending) l = torch.abs(J[:, 13, 1] - J[:, 12, 1]) # Length of the spine section t_posed[:,13] = t_posed[:,13] + \ (curve_torch_3d(head_bending, head_extension, t=interp_t, l=l) - curve_torch_3d(angle_zero, angle_zero, t=interp_t, l=l)) # ------- Body surface transformation matrix---------- G_ = torch.cat([R, t_posed], dim=-1) # BxJx3x4 local transformation matrix pad_row = torch.FloatTensor([0, 0, 0, 1]).to(device).view(1, 1, 1, 4).expand(B, Nj, -1, -1) # BxJx1x4 G_ = torch.cat([G_, pad_row], dim=2) # BxJx4x4 padded to be 4x4 matrix an enable multiplication for the kinematic chain # Global transform G = [G_[:, 0].clone()] for i in range(1, Nj): G.append(torch.matmul(G[self.parent[i - 1]], G_[:, i, :, :])) G = torch.stack(G, dim=1) # ------- Pose dependant blend shapes ---------- if pose_dep_bs is False: v_shaped_pd = v_shaped else: # Note : Those should be retrained for SKEL as the SKEL joints location are different from SMPL. # But the current version lets use get decent pose dependant deformations for the shoulders, belly and knies ident = torch.eye(3, dtype=v_shaped.dtype, device=device) # We need the per SMPL joint bone transform to compute pose dependant blend shapes. # Initialize each joint rotation with identity Rsmpl = ident.unsqueeze(0).unsqueeze(0).expand(B, self.num_joints_smpl, -1, -1).clone() # BxNjx3x3 Rskin = G_[:, :, :3, :3] # BxNjx3x3 Rsmpl[:, smpl_joint_corresp] = Rskin[:] # BxNjx3x3 pose params to rotation pose_feature = Rsmpl[:, 1:].view(B, -1, 3, 3) - ident pose_offsets = torch.matmul(pose_feature.view(B, -1), self.posedirs.view(Ns*3, -1).T).view(B, -1, 3) v_shaped_pd = v_shaped + pose_offsets ########################################################################################## #Transform skin mesh ############################################################################################ # Apply global transformation to the template mesh rest = torch.cat([J, torch.zeros(B, Nj, 1).to(device)], dim=2).view(B, Nj, 4, 1) # BxJx4x1 zeros = torch.zeros(B, Nj, 4, 3).to(device) # BxJx4x3 rest = torch.cat([zeros, rest], dim=-1) # BxJx4x4 rest = torch.matmul(G, rest) # This is a 4x4 transformation matrix that only contains translation to the rest pose joint location Gskin = G - rest # Compute per vertex transformation matrix (after weighting) T = torch.matmul(self.skin_weights, Gskin.permute(1, 0, 2, 3).contiguous().view(Nj, -1)).view(Ns, B, 4,4).transpose(0, 1) rest_shape_h = torch.cat([v_shaped_pd, torch.ones_like(v_shaped_pd)[:, :, [0]]], dim=-1) v_posed = torch.matmul(T, rest_shape_h[:, :, :, None])[:, :, :3, 0] # translation v_trans = v_posed + trans[:,None,:] ########################################################################################## #Transform joints ############################################################################################ # import ipdb; ipdb.set_trace() root_transform = with_zeros(torch.cat((R[:,0],J[:,0][:,:,None]),2)) results = [root_transform] for i in range(0, self.parent.shape[0]): transform_i = with_zeros(torch.cat((R[:, i + 1], t_posed[:,i+1]), 2)) curr_res = torch.matmul(results[self.parent[i]],transform_i) results.append(curr_res) results = torch.stack(results, dim=1) posed_joints = results[:, :, :3, 3] J_transformed = posed_joints + trans[:,None,:] ########################################################################################## # Transform skeleton ############################################################################################ if skelmesh: G_bones = None # Shape the skeleton by scaling its bones skel_rest_shape_h = torch.cat([skel_v0, torch.ones_like(skel_v0)[:, :, [0]]], dim=-1).expand(B, Nk, -1) # (1,Nk,3) # compute the bones scaling from the kinematic tree and skin mesh #with torch.no_grad(): # TODO: when dJ is optimized the shape of the mesh should be affected by the gradients bone_scale = self.compute_bone_scale(J_, v_shaped, skin_v0) # Apply bone meshes scaling: skel_v_shaped = torch.cat([(torch.matmul(bone_scale[:,:,0], self.skel_weights_rigid.T) * skel_rest_shape_h[:, :, 0])[:, :, None], (torch.matmul(bone_scale[:,:,1], self.skel_weights_rigid.T) * skel_rest_shape_h[:, :, 1])[:, :, None], (torch.matmul(bone_scale[:,:,2], self.skel_weights_rigid.T) * skel_rest_shape_h[:, :, 2])[:, :, None], (torch.ones(B, Nk, 1).to(device)) ], dim=-1) # Align the bones with the proper axis Gk01 = build_homog_matrix(Rk01, J.unsqueeze(-1)) # BxJx4x4 T = torch.matmul(self.skel_weights_rigid, Gk01.permute(1, 0, 2, 3).contiguous().view(Nj, -1)).view(Nk, B, 4,4).transpose(0, 1) #[1, 48757, 3, 3] skel_v_align = torch.matmul(T, skel_v_shaped[:, :, :, None])[:, :, :, 0] # This transfo will be applied with weights, effectively unposing the whole skeleton mesh in each joint frame. # Then, per joint weighted transformation can then be applied G_tpose_to_unposed = build_homog_matrix(torch.eye(3).view(1,1,3,3).expand(B, Nj, 3, 3).to(device), -J.unsqueeze(-1)) # BxJx4x4 G_skel = torch.matmul(G, G_tpose_to_unposed) G_bones = torch.matmul(G, Gk01) T = torch.matmul(self.skel_weights, G_skel.permute(1, 0, 2, 3).contiguous().view(Nj, -1)).view(Nk, B, 4,4).transpose(0, 1) skel_v_posed = torch.matmul(T, skel_v_align[:, :, :, None])[:, :, :3, 0] skel_trans = skel_v_posed + trans[:,None,:] else: skel_trans = skel_v0 Gk01 = build_homog_matrix(Rk01, J.unsqueeze(-1)) # BxJx4x4 G_bones = torch.matmul(G, Gk01) joints = J_transformed skin_verts = v_trans skel_verts = skel_trans joints_ori = G_bones[:,:,:3,:3] if skin_verts.max() > 1e3: import ipdb; ipdb.set_trace() output = SKELOutput(skin_verts=skin_verts, skel_verts=skel_verts, joints=joints, joints_ori=joints_ori, betas=betas, poses=poses, trans = trans, pose_offsets = pose_offsets, joints_tpose = J_tpose, v_shaped = v_shaped,) return output def compute_bone_scale(self, J_, v_shaped, skin_v0): # index [0, 1, 2, 3 4, 5, , ...] # todo add last one, figure out bone scale indices # J_ bone vectors [j0, j1-j0, j2-j0, j3-j0, j4-j1, j5-j2, ...] # norm(J) = length of the bone [j0, j1-j0, j2-j0, j3-j0, j4-j1, j5-j2, ...] # self.joints_sockets [j0, j1-j0, j2-j0, j3-j0, j4-j1, j5-j2, ...] # self.skel_weights [j0, j1, j2, j3, j4, j5, ...] B = J_.shape[0] Nj = J_.shape[1] bone_scale = torch.ones(B, Nj).to(J_.device) # BSM template joints location osim_joints_r = self.apose_rel_transfo[:, :3, 3].view(1, Nj, 3).expand(B, Nj, 3).clone() length_bones_bsm = torch.norm(osim_joints_r, dim=-1).expand(B, -1) length_bones_smpl = torch.norm(J_, dim=-1) # (B, Nj) bone_scale_parent = length_bones_smpl / length_bones_bsm non_leaf_node = (self.child != 0) bone_scale[:,non_leaf_node] = (bone_scale_parent[:,self.child])[:,non_leaf_node] # Ulna should have the same scale as radius bone_scale[:,16] = bone_scale[:,17] bone_scale[:,16] = bone_scale[:,17] bone_scale[:,21] = bone_scale[:,22] bone_scale[:,21] = bone_scale[:,22] # Thorax # Thorax scale is defined by the relative position of the thorax to its child joint, not parent joint as for other bones bone_scale[:, 12] = bone_scale[:, 11] # Lumbars # Lumbar scale is defined by the y relative position of the lumbar joint length_bones_bsm = torch.abs(osim_joints_r[:,11, 1]) length_bones_smpl = torch.abs(J_[:, 11, 1]) # (B, Nj) bone_scale_lumbar = length_bones_smpl / length_bones_bsm bone_scale[:, 11] = bone_scale_lumbar # Expand to 3 dimensions and adjest scaling to avoid skin-skeleton intersection and handle the scaling of leaf body parts (hands, feet) bone_scale = bone_scale.reshape(B, Nj, 1).expand(B, Nj, 3).clone() for (ji, doi, dsi), (v1, v2) in scaling_keypoints.items(): bone_scale[:, ji, doi] = ((v_shaped[:,v1] - v_shaped[:, v2])/ (skin_v0[:,v1] - skin_v0[:, v2]))[:,dsi] # Top over chin #TODO: Add keypoints for feet scaling in scaling_keypoints # Adjust thorax front-back scaling # TODO fix this part v1 = 3027 #thorax back v2 = 3495 #thorax front scale_thorax_up = ((v_shaped[:,v1] - v_shaped[:, v2])/ (skin_v0[:,v1] - skin_v0[:, v2]))[:,2] # good for large people v2 = 3506 #sternum scale_thorax_sternum = ((v_shaped[:,v1] - v_shaped[:, v2])/ (skin_v0[:,v1] - skin_v0[:, v2]))[:,2] # Good for skinny people bone_scale[:, 12, 0] = torch.min(scale_thorax_up, scale_thorax_sternum) # Avoids super expanded ribcage for large people and sternum outside for skinny people #lumbars, adjust width to be same as thorax bone_scale[:, 11, 0] = bone_scale[:, 12, 0] return bone_scale def compute_bone_orientation(self, J, J_): """Compute each bone orientation in T pose """ # method = 'unposed' # method = 'learned' method = 'learn_adjust' B = J_.shape[0] Nj = J_.shape[1] # Create an array of bone vectors the bone meshes should be aligned to. bone_vect = torch.zeros_like(J_) # / torch.norm(J_, dim=-1)[:, :, None] # (B, Nj, 3) bone_vect[:] = J_[:, self.child] # Most bones are aligned between their parent and child joint bone_vect[:,16] = bone_vect[:,16]+bone_vect[:,17] # We want to align the ulna to the segment joint 16 to 18 bone_vect[:,21] = bone_vect[:,21]+bone_vect[:,22] # Same other ulna # TODO Check indices here # bone_vect[:,13] = bone_vect[:,12].clone() bone_vect[:,12] = bone_vect.clone()[:,11].clone() # We want to align the thorax on the thorax-lumbar segment # bone_vect[:,11] = bone_vect[:,0].clone() osim_vect = self.apose_rel_transfo[:, :3, 3].clone().view(1, Nj, 3).expand(B, Nj, 3).clone() osim_vect[:] = osim_vect[:,self.child] osim_vect[:,16] = osim_vect[:,16]+osim_vect[:,17] # We want to align the ulna to the segment joint 16 to 18 osim_vect[:,21] = osim_vect[:,21]+osim_vect[:,22] # We want to align the ulna to the segment joint 16 to 18 # TODO: remove when this has been checked # import matplotlib.pyplot as plt # fig = plt.figure() # ax = fig.add_subplot(111, projection='3d') # ax.plot(osim_vect[:,0,0], osim_vect[:,0,1], osim_vect[:,0,2], color='r') # plt.show() Gk = torch.eye(3, device=J_.device).repeat(B, Nj, 1, 1) if method == 'unposed': return Gk elif method == 'learn_adjust': Gk_learned = self.per_joint_rot.view(1, Nj, 3, 3).expand(B, -1, -1, -1) #load learned rotation osim_vect_corr = torch.matmul(Gk_learned, osim_vect.unsqueeze(-1)).squeeze(-1) Gk[:,:] = rotation_matrix_from_vectors(osim_vect_corr, bone_vect) # set nan to zero # TODO: Check again why the following line was required Gk[torch.isnan(Gk)] = 0 # Gk[:,[18,23]] = Gk[:,[16,21]] # hand has same orientation as ulna # Gk[:,[5,10]] = Gk[:,[4,9]] # toe has same orientation as calcaneus # Gk[:,[0,11,12,13,14,19]] = torch.eye(3, device=J_.device).view(1,3,3).expand(B, 6, 3, 3) # pelvis, torso and shoulder blade orientation does not vary with beta, leave it Gk[:, self.joint_idx_fixed_beta] = torch.eye(3, device=J_.device).view(1,3,3).expand(B, len(self.joint_idx_fixed_beta), 3, 3) # pelvis, torso and shoulder blade orientation should not vary with beta, leave it Gk = torch.matmul(Gk, Gk_learned) elif method == 'learned': """ Apply learned transformation""" Gk = self.per_joint_rot.view(1, Nj, 3, 3).expand(B, -1, -1, -1) else: raise NotImplementedError return Gk