Spaces:
Runtime error
Runtime error
| from __future__ import print_function, division | |
| import argparse | |
| import torch | |
| import os,sys | |
| from os import walk, listdir | |
| from os.path import isfile, join | |
| import numpy as np | |
| import joblib | |
| import smplx | |
| import trimesh | |
| import h5py | |
| from tqdm import tqdm | |
| sys.path.append(os.path.join(os.path.dirname(__file__), "src")) | |
| from smplify import SMPLify3D | |
| import config | |
| # parsing argmument | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--batchSize', type=int, default=1, | |
| help='input batch size') | |
| parser.add_argument('--num_smplify_iters', type=int, default=100, | |
| help='num of smplify iters') | |
| parser.add_argument('--cuda', type=bool, default=False, | |
| help='enables cuda') | |
| parser.add_argument('--gpu_ids', type=int, default=0, | |
| help='choose gpu ids') | |
| parser.add_argument('--num_joints', type=int, default=22, | |
| help='joint number') | |
| parser.add_argument('--joint_category', type=str, default="AMASS", | |
| help='use correspondence') | |
| parser.add_argument('--fix_foot', type=str, default="False", | |
| help='fix foot or not') | |
| parser.add_argument('--data_folder', type=str, default="./demo/demo_data/", | |
| help='data in the folder') | |
| parser.add_argument('--save_folder', type=str, default="./demo/demo_results/", | |
| help='results save folder') | |
| parser.add_argument('--files', type=str, default="test_motion.npy", | |
| help='files use') | |
| opt = parser.parse_args() | |
| print(opt) | |
| # ---load predefined something | |
| device = torch.device("cuda:" + str(opt.gpu_ids) if opt.cuda else "cpu") | |
| print(config.SMPL_MODEL_DIR) | |
| smplmodel = smplx.create(config.SMPL_MODEL_DIR, | |
| model_type="smpl", gender="neutral", ext="pkl", | |
| batch_size=opt.batchSize).to(device) | |
| # ## --- load the mean pose as original ---- | |
| smpl_mean_file = config.SMPL_MEAN_FILE | |
| file = h5py.File(smpl_mean_file, 'r') | |
| init_mean_pose = torch.from_numpy(file['pose'][:]).unsqueeze(0).float() | |
| init_mean_shape = torch.from_numpy(file['shape'][:]).unsqueeze(0).float() | |
| cam_trans_zero = torch.Tensor([0.0, 0.0, 0.0]).to(device) | |
| # | |
| pred_pose = torch.zeros(opt.batchSize, 72).to(device) | |
| pred_betas = torch.zeros(opt.batchSize, 10).to(device) | |
| pred_cam_t = torch.zeros(opt.batchSize, 3).to(device) | |
| keypoints_3d = torch.zeros(opt.batchSize, opt.num_joints, 3).to(device) | |
| # # #-------------initialize SMPLify | |
| smplify = SMPLify3D(smplxmodel=smplmodel, | |
| batch_size=opt.batchSize, | |
| joints_category=opt.joint_category, | |
| num_iters=opt.num_smplify_iters, | |
| device=device) | |
| #print("initialize SMPLify3D done!") | |
| purename = os.path.splitext(opt.files)[0] | |
| # --- load data --- | |
| data = np.load(opt.data_folder + "/" + purename + ".npy") # [nframes, njoints, 3] | |
| dir_save = os.path.join(opt.save_folder, purename) | |
| if not os.path.isdir(dir_save): | |
| os.makedirs(dir_save, exist_ok=True) | |
| # run the whole seqs | |
| num_seqs = data.shape[0] | |
| for idx in tqdm(range(num_seqs)): | |
| #print(idx) | |
| joints3d = data[idx] #*1.2 #scale problem [check first] | |
| keypoints_3d[0, :, :] = torch.Tensor(joints3d).to(device).float() | |
| if idx == 0: | |
| pred_betas[0, :] = init_mean_shape | |
| pred_pose[0, :] = init_mean_pose | |
| pred_cam_t[0, :] = cam_trans_zero | |
| else: | |
| data_param = joblib.load(dir_save + "/" + "%04d"%(idx-1) + ".pkl") | |
| pred_betas[0, :] = torch.from_numpy(data_param['beta']).unsqueeze(0).float() | |
| pred_pose[0, :] = torch.from_numpy(data_param['pose']).unsqueeze(0).float() | |
| pred_cam_t[0, :] = torch.from_numpy(data_param['cam']).unsqueeze(0).float() | |
| if opt.joint_category =="AMASS": | |
| confidence_input = torch.ones(opt.num_joints) | |
| # make sure the foot and ankle | |
| if opt.fix_foot == True: | |
| confidence_input[7] = 1.5 | |
| confidence_input[8] = 1.5 | |
| confidence_input[10] = 1.5 | |
| confidence_input[11] = 1.5 | |
| else: | |
| print("Such category not settle down!") | |
| # ----- from initial to fitting ------- | |
| new_opt_vertices, new_opt_joints, new_opt_pose, new_opt_betas, \ | |
| new_opt_cam_t, new_opt_joint_loss = smplify( | |
| pred_pose.detach(), | |
| pred_betas.detach(), | |
| pred_cam_t.detach(), | |
| keypoints_3d, | |
| conf_3d=confidence_input.to(device), | |
| seq_ind=idx | |
| ) | |
| # # -- save the results to ply--- | |
| outputp = smplmodel(betas=new_opt_betas, global_orient=new_opt_pose[:, :3], body_pose=new_opt_pose[:, 3:], | |
| transl=new_opt_cam_t, return_verts=True) | |
| mesh_p = trimesh.Trimesh(vertices=outputp.vertices.detach().cpu().numpy().squeeze(), faces=smplmodel.faces, process=False) | |
| mesh_p.export(dir_save + "/" + "%04d"%idx + ".ply") | |
| # save the pkl | |
| param = {} | |
| param['beta'] = new_opt_betas.detach().cpu().numpy() | |
| param['pose'] = new_opt_pose.detach().cpu().numpy() | |
| param['cam'] = new_opt_cam_t.detach().cpu().numpy() | |
| joblib.dump(param, dir_save + "/" + "%04d"%idx + ".pkl", compress=3) | |