Spaces:
Runtime error
Runtime error
| import torch | |
| import os, sys | |
| import pickle | |
| sys.path.append(os.path.dirname(__file__)) | |
| from customloss import (camera_fitting_loss_3d, | |
| body_fitting_loss_3d, | |
| ) | |
| from prior import MaxMixturePrior | |
| from SMPLX.visualize_joint2smpl.joints2smpl.src import config | |
| from tqdm import tqdm | |
| def guess_init_3d(model_joints, | |
| j3d, | |
| joints_category="orig"): | |
| """Initialize the camera translation via triangle similarity, by using the torso joints . | |
| :param model_joints: SMPL model with pre joints | |
| :param j3d: 25x3 array of Kinect Joints | |
| :returns: 3D vector corresponding to the estimated camera translation | |
| """ | |
| # get the indexed four | |
| gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder'] | |
| gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints] | |
| if joints_category=="orig": | |
| joints_ind_category = [config.JOINT_MAP[joint] for joint in gt_joints] | |
| elif joints_category=="AMASS": | |
| joints_ind_category = [config.AMASS_JOINT_MAP[joint] for joint in gt_joints] | |
| else: | |
| print("NO SUCH JOINTS CATEGORY!") | |
| sum_init_t = (j3d[:, joints_ind_category] - model_joints[:, gt_joints_ind]).sum(dim=1) | |
| init_t = sum_init_t / 4.0 | |
| return init_t | |
| # SMPLIfy 3D | |
| class SMPLify3D(): | |
| """Implementation of SMPLify, use 3D joints.""" | |
| def __init__(self, | |
| smplxmodel, | |
| step_size=1.0, | |
| num_iters=100, | |
| joints_category="orig", | |
| device=torch.device('cuda:0'), | |
| ): | |
| # Store options | |
| self.device = device | |
| self.step_size = step_size | |
| self.num_iters = num_iters | |
| # GMM pose prior | |
| self.pose_prior = MaxMixturePrior(prior_folder=config.GMM_MODEL_DIR, | |
| num_gaussians=8, | |
| dtype=torch.float32).to(device) | |
| # reLoad SMPL-X model | |
| self.smpl = smplxmodel | |
| self.model_faces = smplxmodel.faces_tensor.view(-1) | |
| # select joint joint_category | |
| self.joints_category = joints_category | |
| if joints_category=="orig": | |
| self.smpl_index = config.full_smpl_idx | |
| self.corr_index = config.full_smpl_idx | |
| elif joints_category=="AMASS": | |
| self.smpl_index = config.amass_smpl_idx | |
| self.corr_index = config.amass_idx | |
| else: | |
| self.smpl_index = None | |
| self.corr_index = None | |
| print("NO SUCH JOINTS CATEGORY!") | |
| # ---- get the man function here ------ | |
| def __call__(self, init_pose, init_betas, init_cam_t, j3d, conf_3d=1.0, seq_ind=0): | |
| """Perform body fitting. | |
| Input: | |
| init_pose: SMPL pose estimate | |
| init_betas: SMPL betas estimate | |
| init_cam_t: Camera translation estimate | |
| j3d: joints 3d aka keypoints | |
| conf_3d: confidence for 3d joints | |
| seq_ind: index of the sequence | |
| Returns: | |
| vertices: Vertices of optimized shape | |
| joints: 3D joints of optimized shape | |
| pose: SMPL pose parameters of optimized shape | |
| betas: SMPL beta parameters of optimized shape | |
| camera_translation: Camera translation | |
| """ | |
| # Split SMPL pose to body pose and global orientation | |
| body_pose = init_pose[:, 3:].detach().clone() | |
| global_orient = init_pose[:, :3].detach().clone() | |
| betas = init_betas.detach().clone() | |
| # use guess 3d to get the initial | |
| smpl_output = self.smpl(global_orient=global_orient, | |
| body_pose=body_pose, | |
| betas=betas) | |
| model_joints = smpl_output.joints | |
| init_cam_t = guess_init_3d(model_joints, j3d, self.joints_category).unsqueeze(1).detach() | |
| camera_translation = init_cam_t.clone() | |
| preserve_pose = init_pose[:, 3:].detach().clone() | |
| # -------------Step 1: Optimize camera translation and body orientation-------- | |
| # Optimize only camera translation and body orientation | |
| body_pose.requires_grad = False | |
| betas.requires_grad = True | |
| global_orient.requires_grad = True | |
| camera_translation.requires_grad = True | |
| camera_opt_params = [betas, global_orient, camera_translation] | |
| camera_optimizer = torch.optim.LBFGS(camera_opt_params, max_iter=10, | |
| lr=self.step_size, line_search_fn='strong_wolfe') | |
| cycle = tqdm(range(10)) | |
| for i in cycle: | |
| def closure(): | |
| camera_optimizer.zero_grad() | |
| smpl_output = self.smpl(global_orient=global_orient, | |
| body_pose=body_pose, | |
| betas=betas) | |
| model_joints = smpl_output.joints | |
| loss = camera_fitting_loss_3d(model_joints, camera_translation, | |
| init_cam_t, j3d, self.joints_category) | |
| loss.backward() | |
| return loss | |
| camera_optimizer.step(closure) | |
| # Fix camera translation after optimizing camera | |
| # --------Step 2: Optimize body joints -------------------------- | |
| # Optimize only the body pose and global orientation of the body | |
| body_pose.requires_grad = True | |
| global_orient.requires_grad = True | |
| camera_translation.requires_grad = True | |
| betas.requires_grad = True | |
| body_opt_params = [body_pose, betas, global_orient, camera_translation] | |
| cycle = tqdm(range(self.num_iters)) | |
| body_optimizer = torch.optim.LBFGS(body_opt_params, max_iter=self.num_iters, | |
| lr=self.step_size, line_search_fn='strong_wolfe') | |
| for i in cycle: | |
| def closure(): | |
| body_optimizer.zero_grad() | |
| smpl_output = self.smpl(global_orient=global_orient, | |
| body_pose=body_pose, | |
| betas=betas) | |
| model_joints = smpl_output.joints | |
| model_vertices = smpl_output.vertices | |
| loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation, | |
| j3d[:, self.corr_index], self.pose_prior, | |
| joints3d_conf=conf_3d, | |
| joint_loss_weight=600.0, | |
| pose_preserve_weight=5.0, | |
| use_collision=False, | |
| model_vertices=model_vertices, model_faces=self.model_faces, | |
| search_tree=None, pen_distance=None, filter_faces=None) | |
| loss.backward() | |
| return loss | |
| body_optimizer.step(closure) | |
| pose = torch.cat([global_orient, body_pose], dim=-1).detach() | |
| return pose | |