Spaces:
Paused
Paused
| import torch | |
| import smplx | |
| from constants import * | |
| from scipy.interpolate import interp1d | |
| from torch import nn, einsum | |
| import pytorch3d as T | |
| class JointsToSMPLX(nn.Module): | |
| def __init__(self, input_dim, output_dim, hidden_dim, **kwargs): | |
| super().__init__() | |
| self.layers = nn.Sequential( | |
| nn.Linear(input_dim, hidden_dim), | |
| nn.BatchNorm1d(hidden_dim), | |
| nn.ReLU(), | |
| # nn.Linear(hidden_dim, hidden_dim), | |
| # nn.BatchNorm1d(hidden_dim), | |
| # nn.ReLU(), | |
| nn.Linear(hidden_dim, hidden_dim), | |
| nn.BatchNorm1d(hidden_dim), | |
| nn.ReLU(), | |
| nn.Linear(hidden_dim, output_dim), | |
| ) | |
| def forward(self, x): | |
| return self.layers(x) | |
| def optimize_smpl(pose_pred, joints, joints_ind, hand_pca=45): | |
| device = joints.device | |
| len = joints.shape[0] | |
| smpl_model = smplx.create('./smpl_models', model_type='smplx', | |
| gender='male', ext='npz', | |
| num_betas=10, | |
| use_pca=False, | |
| create_global_orient=True, | |
| create_body_pose=True, | |
| create_betas=True, | |
| create_left_hand_pose=True, | |
| create_right_hand_pose=True, | |
| create_expression=True, | |
| create_jaw_pose=True, | |
| create_leye_pose=True, | |
| create_reye_pose=True, | |
| create_transl=True, | |
| batch_size=len, | |
| ).to(device) | |
| smpl_model.eval() | |
| # weights = torch.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 100, 100, 100, 100]).reshape(nb_joints, 1).repeat(1, 3).to(device) | |
| joints = joints.reshape(len, -1, 3) + torch.tensor(pelvis_shift).to(device) | |
| pose_input = torch.nn.Parameter(pose_pred.detach(), requires_grad=True) | |
| transl = torch.nn.Parameter(torch.zeros(pose_pred.shape[0], 3).to(device), requires_grad=True) | |
| # left_hand = torch.nn.Parameter(torch.zeros(pose_pred.shape[0], hand_pca).to(device), requires_grad=True) | |
| # right_hand = torch.nn.Parameter(torch.zeros(pose_pred.shape[0], hand_pca).to(device), requires_grad=True) | |
| left_hand = torch.from_numpy(relaxed_hand_pose[:45].reshape(1, -1).repeat(pose_pred.shape[0], axis=0)).to(device) | |
| right_hand = torch.from_numpy(relaxed_hand_pose[45:].reshape(1, -1).repeat(pose_pred.shape[0], axis=0)).to(device) | |
| optimizer = torch.optim.Adam(params=[pose_input, transl], lr=0.05) | |
| loss_fn = nn.MSELoss() | |
| vertices_output = None | |
| for step in range(100): | |
| smpl_output = smpl_model(transl=transl, body_pose=pose_input[:, 3:], global_orient=pose_input[:, :3], return_verts=True, | |
| left_hand_pose=left_hand,# @ left_hand_components[:hand_pca], | |
| right_hand_pose=right_hand,# @ right_hand_components[:hand_pca], | |
| ) | |
| joints_output = smpl_output.joints[:, joints_ind].reshape(len, -1, 3) | |
| vertices_output = smpl_output.vertices[:, ::10].detach().cpu().numpy() | |
| loss = loss_fn(joints[:, :], joints_output[:, :]) | |
| # loss = torch.mean((joints - joints_output) ** 2 * weights) | |
| optimizer.zero_grad() | |
| loss.backward() | |
| optimizer.step() | |
| print(loss.item()) | |
| #left_hand = left_hand @ left_hand_components[:hand_pca] | |
| #right_hand = right_hand @ right_hand_components[:hand_pca] | |
| return pose_input.detach().cpu().numpy(), transl.detach().cpu().numpy(), left_hand.detach().cpu().numpy(), right_hand.detach().cpu().numpy(), vertices_output | |
| def joints_to_smpl(model, joints, joints_ind, interp_s): | |
| joints = interpolate_joints(joints, scale=interp_s) | |
| # joints = interpolate_joints(joints, scale=0.33) | |
| # joints = interpolate_joints(joints, scale=interp_s * 3) | |
| input_len = joints.shape[0] | |
| joints = joints.reshape(input_len, -1, 3) | |
| joints = joints.permute(1, 0, 2) | |
| trans_np = joints[0].detach().cpu().numpy() | |
| joints = joints - joints[0] | |
| joints = joints.permute(1, 0, 2) | |
| joints = joints.reshape(input_len, -1) | |
| pose_pred = model(joints) | |
| pose_pred = pose_pred.reshape(-1, 6) | |
| pose_pred = T.matrix_to_axis_angle(T.rotation_6d_to_matrix(pose_pred)).reshape(input_len, -1) | |
| # pose_pred = pose_pred[:seq_len] | |
| pose_output, transl, left_hand, right_hand, vertices = optimize_smpl(pose_pred, joints, joints_ind) | |
| transl = trans_np - np.array(pelvis_shift) + transl | |
| vertices = vertices + transl.reshape(-1, 1, 3) | |
| return pose_output, transl, left_hand, right_hand, vertices | |
| def interpolate_joints(joints, scale): | |
| if scale == 1: | |
| return joints | |
| device = joints.device | |
| joints = joints.detach().cpu().numpy() | |
| in_len = joints.shape[0] | |
| out_len = int(in_len * scale) | |
| joints = joints.reshape(in_len, -1) | |
| x = np.array(range(in_len)) | |
| xnew = np.linspace(0, in_len - 1, out_len) | |
| f = interp1d(x, joints, axis=0) | |
| joints_new = f(xnew) | |
| joints_new = torch.from_numpy(joints_new).to(device).float() | |
| return joints_new |