Spaces:
Runtime error
Runtime error
| import os.path as osp | |
| import argparse | |
| import numpy as np | |
| import torch | |
| import pyrender | |
| import trimesh | |
| import smplx | |
| from tqdm.auto import tqdm, trange | |
| from pathlib import Path | |
| def main(model_folder, | |
| motion_file, | |
| model_type='smplx', | |
| ext='npz', | |
| gender='neutral', | |
| plot_joints=False, | |
| num_betas=10, | |
| sample_expression=True, | |
| num_expression_coeffs=10, | |
| use_face_contour=False): | |
| # open motion file | |
| motion = np.load(motion_file, allow_pickle=True) | |
| _motion = {} | |
| for k,v in motion.items(): | |
| if isinstance(v, np.ndarray): | |
| print(k, motion[k].shape, motion[k].dtype) | |
| if motion[k].dtype in ("<U7", "<U5", "<U4", "object", "|S7"): | |
| _motion[k] = str(motion[k]) | |
| else: | |
| _motion[k] = torch.from_numpy(motion[k]).float() | |
| else: | |
| print(k, v) | |
| _motion[k] = v | |
| motion = _motion | |
| if "poses" in motion: | |
| motion["global_orient"] = motion["root_orient"] | |
| motion["body_pose"] = motion["pose_body"] # seriously? | |
| motion["left_hand_pose"] = motion["pose_hand"][:,:45] | |
| motion["right_hand_pose"] = motion["pose_hand"][:,45:] | |
| num_betas = len(motion['betas']) | |
| gender = str(motion['gender']) | |
| model = smplx.create(model_folder, model_type=model_type, | |
| gender=gender, use_face_contour=use_face_contour, | |
| num_betas=num_betas, | |
| num_expression_coeffs=num_expression_coeffs, | |
| use_pca=False, | |
| ext=ext) | |
| betas, expression = motion['betas'], None | |
| betas = betas.unsqueeze(0)[:, :model.num_betas] | |
| global_orient = motion['global_orient'] | |
| body_pose = motion['body_pose'] | |
| left_hand_pose = motion['left_hand_pose'] | |
| right_hand_pose = motion['right_hand_pose'] | |
| # if sample_expression: | |
| # expression = torch.randn( | |
| # [1, model.num_expression_coeffs], dtype=torch.float32) | |
| #print(expression) | |
| #print(betas.shape, body_pose.shape, expression.shape) | |
| for pose_idx in trange(body_pose.size(0)): | |
| pose_idx = [pose_idx] | |
| # output = model(betas=betas, # expression=expression, | |
| # return_verts=True) | |
| # for x in [betas, global_orient, body_pose, left_hand_pose, right_hand_pose]: | |
| # print(x.dtype, x.shape) | |
| output = model( | |
| betas=betas, | |
| global_orient=global_orient[pose_idx], | |
| body_pose=body_pose[pose_idx], | |
| left_hand_pose=left_hand_pose[pose_idx], | |
| right_hand_pose=right_hand_pose[pose_idx], | |
| # expression=expression, | |
| return_verts=True | |
| ) | |
| vertices = output.vertices.detach().cpu().numpy().squeeze() | |
| joints = output.joints.detach().cpu().numpy().squeeze() | |
| vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8] | |
| tri_mesh = trimesh.Trimesh(vertices, model.faces, | |
| vertex_colors=vertex_colors) | |
| mesh = pyrender.Mesh.from_trimesh(tri_mesh) | |
| scene = pyrender.Scene() | |
| scene.add(mesh) | |
| if plot_joints: | |
| sm = trimesh.creation.uv_sphere(radius=0.005) | |
| sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0] | |
| tfs = np.tile(np.eye(4), (len(joints), 1, 1)) | |
| tfs[:, :3, 3] = joints | |
| joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs) | |
| scene.add(joints_pcl) | |
| pyrender.Viewer(scene, use_raymond_lighting=True) | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser(description='SMPL-X Demo') | |
| parser.add_argument('--model-folder', required=True, type=str, | |
| help='The path to the model folder') | |
| parser.add_argument('--motion-file', required=True, type=str, | |
| help='The path to the motion file to process') | |
| parser.add_argument('--num-expression-coeffs', default=10, type=int, | |
| dest='num_expression_coeffs', | |
| help='Number of expression coefficients.') | |
| parser.add_argument('--ext', type=str, default='npz', | |
| help='Which extension to use for loading') | |
| parser.add_argument('--sample-expression', default=True, | |
| dest='sample_expression', | |
| type=lambda arg: arg.lower() in ['true', '1'], | |
| help='Sample a random expression') | |
| parser.add_argument('--use-face-contour', default=False, | |
| type=lambda arg: arg.lower() in ['true', '1'], | |
| help='Compute the contour of the face') | |
| args = parser.parse_args() | |
| def resolve(path): | |
| return osp.expanduser(osp.expandvars(path)) | |
| model_folder = resolve(args.model_folder) | |
| motion_file = resolve(args.motion_file) | |
| ext = args.ext | |
| num_expression_coeffs = args.num_expression_coeffs | |
| sample_expression = args.sample_expression | |
| main(model_folder, motion_file, ext=ext, | |
| sample_expression=sample_expression, | |
| use_face_contour=args.use_face_contour) | |