def fit(self, keypoints3d, dtype='coco', verbose=True): """Run fitting to optimize the SMPL parameters.""" assert dtype == 'coco', 'only support coco format for now.' assert len( keypoints3d.shape) == 3, 'input shape should be [N, njoints, 3]' mapping_target = unify_joint_mappings(dataset=dtype) keypoints3d = keypoints3d[:, mapping_target, :] keypoints3d = torch.from_numpy(keypoints3d).float().to(self.device) batch_size, njoints = keypoints3d.shape[0:2] # Init learnable smpl model smpl = SMPL(model_path=self.smpl_model_path, gender=self.smpl_model_gender, batch_size=batch_size).to(self.device) # Start fitting for step in range(self.niter): optimizer = self.get_optimizer(smpl, step, self.base_lr) output = smpl.forward() joints = output.joints[:, self.joints_mapping_smpl[:njoints], :] loss = self.metric(joints, keypoints3d) optimizer.zero_grad() loss.backward() optimizer.step() if verbose and step % 10 == 0: logging.info(f'step {step:03d}; loss {loss.item():.3f};') # Return results return smpl, loss.item()
def compute_pve_neutral_pose_scale_corrected(predicted_smpl_shape, target_smpl_shape, gender): """ Given predicted and target SMPL shape parameters, computes neutral-pose per-vertex error after scale-correction (to account for scale vs camera depth ambiguity). :param predicted_smpl_parameters: predicted SMPL shape parameters tensor with shape (1, 10) :param target_smpl_parameters: target SMPL shape parameters tensor with shape (1, 10) :param gender: gender of target """ smpl_male = SMPL(config.SMPL_MODEL_DIR, batch_size=1, gender='male') smpl_female = SMPL(config.SMPL_MODEL_DIR, batch_size=1, gender='female') # Get neutral pose vertices if gender == 'm': pred_smpl_neutral_pose_output = smpl_male(betas=predicted_smpl_shape) target_smpl_neutral_pose_output = smpl_male(betas=target_smpl_shape) elif gender == 'f': pred_smpl_neutral_pose_output = smpl_female(betas=predicted_smpl_shape) target_smpl_neutral_pose_output = smpl_female(betas=target_smpl_shape) pred_smpl_neutral_pose_vertices = pred_smpl_neutral_pose_output.vertices target_smpl_neutral_pose_vertices = target_smpl_neutral_pose_output.vertices # Rescale such that RMSD of predicted vertex mesh is the same as RMSD of target mesh. # This is done to combat scale vs camera depth ambiguity. pred_smpl_neutral_pose_vertices_rescale = scale_and_translation_transform_batch( pred_smpl_neutral_pose_vertices, target_smpl_neutral_pose_vertices) # Compute PVE-T-SC pve_neutral_pose_scale_corrected = np.linalg.norm( pred_smpl_neutral_pose_vertices_rescale - target_smpl_neutral_pose_vertices, axis=-1) # (1, 6890) return pve_neutral_pose_scale_corrected
def __init__(self, smpl_path, joints_regressor): super().__init__() assert has_smpl, 'Please install smplx to use SMPL.' self.smpl_neutral = SMPL_( model_path=smpl_path, create_global_orient=False, create_body_pose=False, create_transl=False, gender='neutral') self.smpl_male = SMPL_( model_path=smpl_path, create_betas=False, create_global_orient=False, create_body_pose=False, create_transl=False, gender='male') self.smpl_female = SMPL_( model_path=smpl_path, create_betas=False, create_global_orient=False, create_body_pose=False, create_transl=False, gender='female') joints_regressor = torch.tensor( np.load(joints_regressor), dtype=torch.float)[None, ...] self.register_buffer('joints_regressor', joints_regressor) self.num_verts = self.smpl_neutral.get_num_verts() self.num_joints = self.joints_regressor.shape[1]
def main(_): # Parsing data info. aist_dataset = AISTDataset(FLAGS.anno_dir) video_path = os.path.join(FLAGS.video_dir, f'{FLAGS.video_name}.mp4') seq_name, view = AISTDataset.get_seq_name(FLAGS.video_name) view_idx = AISTDataset.VIEWS.index(view) # Parsing keypoints. if FLAGS.mode == '2D': # raw keypoints detection results. keypoints2d, _, _ = AISTDataset.load_keypoint2d( aist_dataset.keypoint2d_dir, seq_name) keypoints2d = keypoints2d[view_idx, :, :, 0:2] elif FLAGS.mode == '3D': # 3D keypoints with temporal optimization. keypoints3d = AISTDataset.load_keypoint3d( aist_dataset.keypoint3d_dir, seq_name, use_optim=True) nframes, njoints, _ = keypoints3d.shape env_name = aist_dataset.mapping_seq2env[seq_name] cgroup = AISTDataset.load_camera_group(aist_dataset.camera_dir, env_name) keypoints2d = cgroup.project(keypoints3d) keypoints2d = keypoints2d.reshape(9, nframes, njoints, 2)[view_idx] elif FLAGS.mode == 'SMPL': # SMPL joints smpl_poses, smpl_scaling, smpl_trans = AISTDataset.load_motion( aist_dataset.motion_dir, seq_name) smpl = SMPL(model_path=FLAGS.smpl_dir, gender='MALE', batch_size=1) keypoints3d = smpl.forward( global_orient=torch.from_numpy(smpl_poses[:, 0:1]).float(), body_pose=torch.from_numpy(smpl_poses[:, 1:]).float(), transl=torch.from_numpy(smpl_trans).float(), scaling=torch.from_numpy(smpl_scaling.reshape(1, 1)).float(), ).joints.detach().numpy() nframes, njoints, _ = keypoints3d.shape env_name = aist_dataset.mapping_seq2env[seq_name] cgroup = AISTDataset.load_camera_group(aist_dataset.camera_dir, env_name) keypoints2d = cgroup.project(keypoints3d) keypoints2d = keypoints2d.reshape(9, nframes, njoints, 2)[view_idx] # Visualize. os.makedirs(FLAGS.save_dir, exist_ok=True) save_path = os.path.join(FLAGS.save_dir, f'{FLAGS.video_name}.mp4') plot_on_video(keypoints2d, video_path, save_path, fps=60)
def fit(self, keypoints3d, dtype='coco', verbose=True): """Run fitting to optimize the SMPL parameters.""" assert dtype == 'coco', 'only support coco format for now.' assert len( keypoints3d.shape) == 3, 'input shape should be [N, njoints, 3]' mapping_target = unify_joint_mappings(dataset=dtype) keypoints3d = keypoints3d[:, mapping_target, :] keypoints3d = torch.from_numpy(keypoints3d).float().to(self.device) batch_size, njoints = keypoints3d.shape[0:2] # Init learnable smpl model smpl = SMPL(model_path=self.smpl_model_path, gender=self.smpl_model_gender, batch_size=batch_size).to(self.device) # Start fitting for step in range(self.niter): optimizer = self.get_optimizer(smpl, step, self.base_lr) output = smpl.forward() joints = output.joints[:, self.joints_mapping_smpl[:njoints], :] loss = self.metric(joints, keypoints3d) optimizer.zero_grad() loss.backward() optimizer.step() if verbose and step % 10 == 0: logging.info(f'step {step:03d}; loss {loss.item():.3f};') if FLAGS.visualize: vertices = output.vertices[0].detach().cpu().numpy( ) # first frame mesh = trimesh.Trimesh(vertices, smpl.faces) mesh.visual.face_colors = [200, 200, 250, 100] pts = vedo.Points(keypoints3d[0].detach().cpu().numpy(), r=20) # first frame vedo.show(mesh, pts, interactive=False) # Return results return smpl, loss.item()
def load_dance_data(dance_dir): print('---------- Loading pose keypoints ----------') aist_dataset = AISTDataset(dance_dir) seq_names = list(aist_dataset.mapping_seq2env.keys()) print(seq_names) dances = {} for seq_name in tqdm(seq_names): print(f'Process -> {seq_name}') smpl_poses, smpl_scaling, smpl_trans = AISTDataset.load_motion( aist_dataset.motion_dir, seq_name) smpl = SMPL(model_path=args.smpl_dir, gender='MALE', batch_size=1) keypoints3d = smpl.forward( global_orient=torch.from_numpy(smpl_poses[:, 0:1]).float(), body_pose=torch.from_numpy(smpl_poses[:, 1:]).float(), transl=torch.from_numpy(smpl_trans / smpl_scaling).float(), ).joints.detach().numpy()[:, 0:24, :] nframes = keypoints3d.shape[0] dances[seq_name] = keypoints3d.reshape(nframes, -1).tolist() print(np.shape(dances[seq_name])) # (nframes, 72) return dances
def get_verts(theta, smpl_path): device = 'cpu' smpl = SMPL(smpl_path, batch_size=1).to(device) pose, betas = theta[:, :72], theta[:, 72:] verts = [] b_ = torch.split(betas, 500) p_ = torch.split(pose, 500) for b, p in zip(b_, p_): output = smpl(betas=b, body_pose=p[:, 3:], global_orient=p[:, :3], pose2rot=True) verts.append(output.vertices.detach().cpu().numpy()) verts = np.concatenate(verts, axis=0) del smpl return verts
def __init__(self, backbone, mesh_head, smpl, disc=None, loss_gan=None, loss_mesh=None, train_cfg=None, test_cfg=None, pretrained=None): super().__init__() assert has_smpl, 'Please install smplx to use SMPL.' self.backbone = builder.build_backbone(backbone) self.mesh_head = builder.build_head(mesh_head) self.generator = torch.nn.Sequential(self.backbone, self.mesh_head) self.smpl = SMPL( model_path=smpl['smpl_path'], create_betas=False, create_global_orient=False, create_body_pose=False, create_transl=False) joints_regressor = torch.tensor( np.load(smpl['joints_regressor']), dtype=torch.float).unsqueeze(0) self.register_buffer('joints_regressor', joints_regressor) self.with_gan = disc is not None and loss_gan is not None if self.with_gan: self.discriminator = SMPLDiscriminator(**disc) self.loss_gan = builder.build_loss(loss_gan) self.disc_step_count = 0 self.train_cfg = train_cfg self.test_cfg = test_cfg self.loss_mesh = builder.build_loss(loss_mesh) self.init_weights(pretrained=pretrained)
def visEFT_singleSubject(renderer): inputDir = args.fit_dir imgDir = args.img_dir smplModelDir = args.smpl_dir smpl = SMPL(smplModelDir, batch_size=1, create_transl=False) # outputFolder = os.path.basename(inputDir) + '_dpOut' # outputFolder =os.path.join('/run/media/hjoo/disk/data/eftout/',outputFolder) eft_fileList = listdir(inputDir) #Check all fitting files print(">> Found {} files in the fitting folder {}".format(len(eft_fileList), inputDir)) totalCnt =0 erroneousCnt =0 for idx, f in enumerate(sorted(eft_fileList)): #Load EFT data fileFullPath = join(inputDir, f) with open(fileFullPath,'rb') as f: eft_data = pickle.load(f) #Get raw image path imgFullPath = eft_data['imageName'][0] imgName = os.path.basename(imgFullPath) imgFullPath =os.path.join(imgDir, os.path.basename(imgFullPath) ) assert os.path.exists(imgFullPath) rawImg = cv2.imread(imgFullPath) print(f'Input image: {imgFullPath}') #EFT data bbox_scale = eft_data['scale'][0] bbox_center = eft_data['center'][0] pred_camera = eft_data['pred_camera'] pred_betas = torch.from_numpy(eft_data['pred_shape']) pred_pose_rotmat = torch.from_numpy(eft_data['pred_pose_rotmat']) #COCO only. Annotation index print("COCO annotId: {}".format(eft_data['annotId'])) #Obtain skeleton and smpl data smpl_output = smpl(betas=pred_betas, body_pose=pred_pose_rotmat[:,1:], global_orient=pred_pose_rotmat[:,0].unsqueeze(1), pose2rot=False ) smpl_vertices = smpl_output.vertices.detach().cpu().numpy() smpl_joints_3d = smpl_output.joints.detach().cpu().numpy() #Crop image croppedImg, boxScale_o2n, bboxTopLeft = crop_bboxInfo(rawImg, bbox_center, bbox_scale, (BBOX_IMG_RES, BBOX_IMG_RES) ) ######################## # Visualize if False: #Compute 2D reprojection error # if not (data['loss_keypoints_2d']<0.0001 or data['loss_keypoints_2d']>0.001 : # continue maxBeta = abs(torch.max( abs(pred_betas)).item()) if eft_data['loss_keypoints_2d']>0.0005 or maxBeta>3: erroneousCnt +=1 print(">>> loss2d: {}, maxBeta: {}".format( eft_data['loss_keypoints_2d'],maxBeta) ) # Visualize 2D image if False: viewer2D.ImShow(rawImg, name='rawImg', waitTime=1) #You should press any key viewer2D.ImShow(croppedImg, name='croppedImg', waitTime=1) # Visualization Mesh if True: b=0 camParam_scale = pred_camera[b,0] camParam_trans = pred_camera[b,1:] pred_vert_vis = smpl_vertices[b] smpl_joints_3d_vis = smpl_joints_3d[b] if args.onbbox: pred_vert_vis = convert_smpl_to_bbox(pred_vert_vis, camParam_scale, camParam_trans) smpl_joints_3d_vis = convert_smpl_to_bbox(smpl_joints_3d_vis, camParam_scale, camParam_trans) renderer.setBackgroundTexture(croppedImg) renderer.setViewportSize(croppedImg.shape[1], croppedImg.shape[0]) else: #Covert SMPL to BBox first pred_vert_vis = convert_smpl_to_bbox(pred_vert_vis, camParam_scale, camParam_trans) smpl_joints_3d_vis = convert_smpl_to_bbox(smpl_joints_3d_vis, camParam_scale, camParam_trans) #From cropped space to original pred_vert_vis = convert_bbox_to_oriIm(pred_vert_vis, boxScale_o2n, bboxTopLeft, rawImg.shape[1], rawImg.shape[0]) smpl_joints_3d_vis = convert_bbox_to_oriIm(smpl_joints_3d_vis, boxScale_o2n, bboxTopLeft, rawImg.shape[1], rawImg.shape[0]) renderer.setBackgroundTexture(rawImg) renderer.setViewportSize(rawImg.shape[1], rawImg.shape[0]) pred_meshes = {'ver': pred_vert_vis, 'f': smpl.faces} v = pred_meshes['ver'] f = pred_meshes['f'] #Visualize in the original image space renderer.set_mesh(v,f) renderer.showBackground(True) renderer.setWorldCenterBySceneCenter() renderer.setCameraViewMode("cam") renderer.setViewportSize(rawImg.shape[1], rawImg.shape[0]) renderer.display() renderImg = renderer.get_screen_color_ibgr() viewer2D.ImShow(renderImg,waitTime=1) # out_all_f = render.get_z_value() # Visualization Mesh on side view if True: # renderer.set_viewpoint() renderer.showBackground(False) renderer.setWorldCenterBySceneCenter() renderer.setCameraViewMode("side") renderer.setViewportSize(rawImg.shape[1], rawImg.shape[0]) renderer.display() sideImg = renderer.get_screen_color_ibgr() #Overwite on rawImg viewer2D.ImShow(sideImg,waitTime=1) sideImg = cv2.resize(sideImg, (renderImg.shape[1], renderImg.shape[0]) ) saveImg = np.concatenate( (renderImg,sideImg), axis =1) viewer2D.ImShow(saveImg,waitTime=0) if True: #Save the rendered image to files if os.path.exists(args.render_dir) == False: os.mkdir(args.render_dir) render_output_path = args.render_dir + '/render_{:08d}.jpg'.format(idx) print(f"Save to {render_output_path}") cv2.imwrite(render_output_path, saveImg) print("erroneous Num : {}/{} ({} percent)".format(erroneousCnt,totalCnt, float(erroneousCnt)*100/totalCnt))
def visEFT_multiSubjects(renderer): inputDir = args.fit_dir imgDir = args.img_dir smplModelDir = args.smpl_dir smpl = SMPL(smplModelDir, batch_size=1, create_transl=False) eft_fileList = listdir(inputDir) #Check all fitting files print(">> Found {} files in the fitting folder {}".format(len(eft_fileList), inputDir)) #Aggregate all efl per image eft_perimage ={} for f in sorted(eft_fileList): #Load imageName = f[:f.rfind('_')] if imageName not in eft_perimage.keys(): eft_perimage[imageName] =[] eft_perimage[imageName].append(f) for imgName in eft_perimage: eftFiles_perimage = eft_perimage[imgName] renderer.clear_mesh() for idx,f in enumerate(eftFiles_perimage): #Load EFT data fileFullPath = join(inputDir, f) with open(fileFullPath,'rb') as f: eft_data = pickle.load(f) #Get raw image path if idx==0: imgFullPath = eft_data['imageName'][0] imgFullPath =os.path.join(imgDir, os.path.basename(imgFullPath) ) assert os.path.exists(imgFullPath) rawImg = cv2.imread(imgFullPath) print(f'Input image: {imgFullPath}') #EFT data bbox_scale = eft_data['scale'][0] bbox_center = eft_data['center'][0] pred_camera = eft_data['pred_camera'] pred_betas = torch.from_numpy(eft_data['pred_shape']) pred_pose_rotmat = torch.from_numpy(eft_data['pred_pose_rotmat']) #COCO only. Annotation index print("COCO annotId: {}".format(eft_data['annotId'])) #Obtain skeleton and smpl data smpl_output = smpl(betas=pred_betas, body_pose=pred_pose_rotmat[:,1:], global_orient=pred_pose_rotmat[:,0].unsqueeze(1), pose2rot=False ) smpl_vertices = smpl_output.vertices.detach().cpu().numpy() smpl_joints_3d = smpl_output.joints.detach().cpu().numpy() #Crop image croppedImg, boxScale_o2n, bboxTopLeft = crop_bboxInfo(rawImg.copy(), bbox_center, bbox_scale, (BBOX_IMG_RES, BBOX_IMG_RES) ) ######################## # Visualize # Visualize 2D image if False: viewer2D.ImShow(rawImg, name='rawImg', waitTime=1) #You should press any key viewer2D.ImShow(croppedImg, name='croppedImg', waitTime=0) # Visualization Mesh on raw images if True: b=0 camParam_scale = pred_camera[b,0] camParam_trans = pred_camera[b,1:] pred_vert_vis = smpl_vertices[b] smpl_joints_3d_vis = smpl_joints_3d[b] if False:#args.onbbox: #Always in the original image pred_vert_vis = convert_smpl_to_bbox(pred_vert_vis, camParam_scale, camParam_trans) smpl_joints_3d_vis = convert_smpl_to_bbox(smpl_joints_3d_vis, camParam_scale, camParam_trans) renderer.setBackgroundTexture(croppedImg) renderer.setViewportSize(croppedImg.shape[1], croppedImg.shape[0]) else: #Covert SMPL to BBox first pred_vert_vis = convert_smpl_to_bbox(pred_vert_vis, camParam_scale, camParam_trans) smpl_joints_3d_vis = convert_smpl_to_bbox(smpl_joints_3d_vis, camParam_scale, camParam_trans) #From cropped space to original pred_vert_vis = convert_bbox_to_oriIm(pred_vert_vis, boxScale_o2n, bboxTopLeft, rawImg.shape[1], rawImg.shape[0]) smpl_joints_3d_vis = convert_bbox_to_oriIm(smpl_joints_3d_vis, boxScale_o2n, bboxTopLeft, rawImg.shape[1], rawImg.shape[0]) renderer.setBackgroundTexture(rawImg) renderer.setViewportSize(rawImg.shape[1], rawImg.shape[0]) pred_meshes = {'ver': pred_vert_vis, 'f': smpl.faces} v = pred_meshes['ver'] f = pred_meshes['f'] #Visualize in the original image spaceq # renderer.set_mesh(v,f) renderer.add_mesh(v,f) #Render Mesh on the camera view renderer.showBackground(True) renderer.setWorldCenterBySceneCenter() renderer.setCameraViewMode("cam") renderer.display() overlaid = renderer.get_screen_color_ibgr() #Overwite on rawImg viewer2D.ImShow(overlaid,waitTime=1,name="overlaid") #Render Mesh on the rotating view renderer.showBackground(False) renderer.setWorldCenterBySceneCenter() renderer.setCameraViewMode("free") for i in range(90): renderer.setViewAngle(i*4,0) renderer.display() sideImg = renderer.get_screen_color_ibgr() #Overwite on rawImg viewer2D.ImShow(sideImg,waitTime=1,name="otherviews") if True: #Save the rendered image to files if os.path.exists(args.render_dir) == False: os.mkdir(args.render_dir) render_output_path = args.render_dir + '/render_{}.jpg'.format(imgName) print(f"Save to {render_output_path}") cv2.imwrite(render_output_path, rawImg)
def __init__(self, focal_length=1000, height=512, width=512): self.renderer = pyrender.OffscreenRenderer(height, width) smpl = SMPL('data/smpl') self.faces = smpl.faces self.focal_length = focal_length
class SMPL(nn.Module): """SMPL 3d human mesh model of paper ref: Matthew Loper. ``SMPL: A skinned multi-person linear model''. This module is based on the smplx project (https://github.com/vchoutas/smplx). Args: smpl_path (str): The path to the folder where the model weights are stored. joints_regressor (str): The path to the file where the joints regressor weight are stored. """ def __init__(self, smpl_path, joints_regressor): super().__init__() assert has_smpl, 'Please install smplx to use SMPL.' self.smpl_neutral = SMPL_( model_path=smpl_path, create_global_orient=False, create_body_pose=False, create_transl=False, gender='neutral') self.smpl_male = SMPL_( model_path=smpl_path, create_betas=False, create_global_orient=False, create_body_pose=False, create_transl=False, gender='male') self.smpl_female = SMPL_( model_path=smpl_path, create_betas=False, create_global_orient=False, create_body_pose=False, create_transl=False, gender='female') joints_regressor = torch.tensor( np.load(joints_regressor), dtype=torch.float)[None, ...] self.register_buffer('joints_regressor', joints_regressor) self.num_verts = self.smpl_neutral.get_num_verts() self.num_joints = self.joints_regressor.shape[1] def smpl_forward(self, model, **kwargs): """Apply a specific SMPL model with given model parameters. Note: B: batch size V: number of vertices K: number of joints Returns: outputs (dict): Dict with mesh vertices and joints. - vertices: Tensor([B, V, 3]), mesh vertices - joints: Tensor([B, K, 3]), 3d joints regressed from mesh vertices. """ betas = kwargs['betas'] batch_size = betas.shape[0] device = betas.device output = {} if batch_size == 0: output['vertices'] = betas.new_zeros([0, self.num_verts, 3]) output['joints'] = betas.new_zeros([0, self.num_joints, 3]) else: smpl_out = model(**kwargs) output['vertices'] = smpl_out.vertices output['joints'] = torch.matmul( self.joints_regressor.to(device), output['vertices']) return output def get_faces(self): """Return mesh faces. Note: F: number of faces Returns: faces: np.ndarray([F, 3]), mesh faces """ return self.smpl_neutral.faces def forward(self, betas, body_pose, global_orient, transl=None, gender=None): """Forward function. Note: B: batch size J: number of controllable joints of model, for smpl model J=23 K: number of joints Args: betas: Tensor([B, 10]), human body shape parameters of SMPL model. body_pose: Tensor([B, J*3] or [B, J, 3, 3]), human body pose parameters of SMPL model. It should be axis-angle vector ([B, J*3]) or rotation matrix ([B, J, 3, 3)]. global_orient: Tensor([B, 3] or [B, 1, 3, 3]), global orientation of human body. It should be axis-angle vector ([B, 3]) or rotation matrix ([B, 1, 3, 3)]. transl: Tensor([B, 3]), global translation of human body. gender: Tensor([B]), gender parameters of human body. -1 for neutral, 0 for male , 1 for female. Returns: outputs (dict): Dict with mesh vertices and joints. - vertices: Tensor([B, V, 3]), mesh vertices - joints: Tensor([B, K, 3]), 3d joints regressed from mesh vertices. """ batch_size = betas.shape[0] pose2rot = True if body_pose.dim() == 2 else False if batch_size > 0 and gender is not None: output = { 'vertices': betas.new_zeros([batch_size, self.num_verts, 3]), 'joints': betas.new_zeros([batch_size, self.num_joints, 3]) } mask = gender < 0 _out = self.smpl_forward( self.smpl_neutral, betas=betas[mask], body_pose=body_pose[mask], global_orient=global_orient[mask], transl=transl[mask] if transl is not None else None, pose2rot=pose2rot) output['vertices'][mask] = _out['vertices'] output['joints'][mask] = _out['joints'] mask = gender == 0 _out = self.smpl_forward( self.smpl_male, betas=betas[mask], body_pose=body_pose[mask], global_orient=global_orient[mask], transl=transl[mask] if transl is not None else None, pose2rot=pose2rot) output['vertices'][mask] = _out['vertices'] output['joints'][mask] = _out['joints'] mask = gender == 1 _out = self.smpl_forward( self.smpl_male, betas=betas[mask], body_pose=body_pose[mask], global_orient=global_orient[mask], transl=transl[mask] if transl is not None else None, pose2rot=pose2rot) output['vertices'][mask] = _out['vertices'] output['joints'][mask] = _out['joints'] else: return self.smpl_forward( self.smpl_neutral, betas=betas, body_pose=body_pose, global_orient=global_orient, transl=transl, pose2rot=pose2rot) return output
import torch # import matplotlib # matplotlib.use('MACOSX') import matplotlib.pyplot as plt from smplx import SMPL from utils.renderer import Renderer from utils.cam_utils import perspective_project_torch from data.ssp3d_dataset import SSP3DDataset import config # SMPL models in torch smpl_male = SMPL(config.SMPL_MODEL_DIR, batch_size=1, gender='male') smpl_female = SMPL(config.SMPL_MODEL_DIR, batch_size=1, gender='female') # Pyrender renderer renderer = Renderer(faces=smpl_male.faces, img_res=512) # SSP-3D datset class ssp3d_dataset = SSP3DDataset(config.SSP_3D_PATH) indices_to_plot = [11, 60, 199] # Visualising 3 examples from SSP-3D for i in indices_to_plot: data = ssp3d_dataset.__getitem__(i) fname = data['fname'] image = data['image'] cropped_image = data['cropped_image'] silhouette = data['silhouette'] joints2D = data['joints2D']
def main(): args = get_args() if args.aist: import vedo if not os.path.exists(args.json_dir): os.makedirs(args.json_dir) if args.aist: print ("test with AIST++ dataset!") music_data, dance_data, dance_names = load_data_aist( args.input_dir, interval=None, rotmat=args.rotmat) else: music_data, dance_data, dance_names = load_data( args.input_dir, interval=None) device = torch.device('cuda' if args.cuda else 'cpu') test_loader = torch.utils.data.DataLoader( DanceDataset(music_data, dance_data), batch_size=args.batch_size, collate_fn=paired_collate_fn ) generator = Generator(args.model, device) if args.aist and args.rotmat: from smplx import SMPL smpl = SMPL(model_path="/media/ruilongli/hd1/Data/smpl/", gender='MALE', batch_size=1) results = [] random_id = 0 # np.random.randint(0, 1e4) for i, batch in enumerate(tqdm(test_loader, desc='Generating dance poses')): # Prepare data src_seq, src_pos, tgt_pose = map(lambda x: x.to(device), batch) pose_seq = generator.generate(src_seq[:, :1200], src_pos[:, :1200]) # first 20 secs results.append(pose_seq) if args.aist: np_dance = pose_seq[0].data.cpu().numpy() if args.rotmat: root = np_dance[:, :3] rotmat = np_dance[:, 3:].reshape([-1, 3, 3]) rotmat = get_closest_rotmat(rotmat) smpl_poses = rotmat2aa(rotmat).reshape(-1, 24, 3) np_dance = smpl.forward( global_orient=torch.from_numpy(smpl_poses[:, 0:1]).float(), body_pose=torch.from_numpy(smpl_poses[:, 1:]).float(), transl=torch.from_numpy(root).float(), ).joints.detach().numpy()[:, 0:24, :] else: root = np_dance[:, :3] np_dance = np_dance + np.tile(root, (1, 24)) np_dance[:, :3] = root np_dance = np_dance.reshape(np_dance.shape[0], -1, 3) print (np_dance.shape) # save save_path = os.path.join(args.json_dir, dance_names[i]+f"_{random_id:04d}") np.save(save_path, np_dance) # visualize for frame in np_dance: pts = vedo.Points(frame, r=20) vedo.show(pts, interactive=False) time.sleep(0.1) exit() if args.aist: pass else: # Visualize generated dance poses np_dances = [] for i in range(len(results)): np_dance = results[i][0].data.cpu().numpy() root = np_dance[:, 2*8:2*9] np_dance = np_dance + np.tile(root, (1, 25)) np_dance[:, 2*8:2*9] = root np_dances.append(np_dance) write2json(np_dances, dance_names, args) visualize(args)
def main(_): # Parsing data info. aist_dataset = AISTDataset(FLAGS.anno_dir) video_path = os.path.join(FLAGS.video_dir, f'{FLAGS.video_name}.mp4') seq_name, view = AISTDataset.get_seq_name(FLAGS.video_name) view_idx = AISTDataset.VIEWS.index(view) # Parsing keypoints. if FLAGS.mode == '2D': # raw keypoints detection results. keypoints2d, _, _ = AISTDataset.load_keypoint2d( aist_dataset.keypoint2d_dir, seq_name) keypoints2d = keypoints2d[view_idx, :, :, 0:2] elif FLAGS.mode == '3D': # 3D keypoints with temporal optimization. keypoints3d = AISTDataset.load_keypoint3d(aist_dataset.keypoint3d_dir, seq_name, use_optim=True) nframes, njoints, _ = keypoints3d.shape env_name = aist_dataset.mapping_seq2env[seq_name] cgroup = AISTDataset.load_camera_group(aist_dataset.camera_dir, env_name) keypoints2d = cgroup.project(keypoints3d) keypoints2d = keypoints2d.reshape(9, nframes, njoints, 2)[view_idx] elif FLAGS.mode == 'SMPL': # SMPL joints smpl_poses, smpl_scaling, smpl_trans = AISTDataset.load_motion( aist_dataset.motion_dir, seq_name) smpl = SMPL(model_path=FLAGS.smpl_dir, gender='MALE', batch_size=1) keypoints3d = smpl.forward( global_orient=torch.from_numpy(smpl_poses[:, 0:1]).float(), body_pose=torch.from_numpy(smpl_poses[:, 1:]).float(), transl=torch.from_numpy(smpl_trans).float(), scaling=torch.from_numpy(smpl_scaling.reshape(1, 1)).float(), ).joints.detach().numpy() nframes, njoints, _ = keypoints3d.shape env_name = aist_dataset.mapping_seq2env[seq_name] cgroup = AISTDataset.load_camera_group(aist_dataset.camera_dir, env_name) keypoints2d = cgroup.project(keypoints3d) keypoints2d = keypoints2d.reshape(9, nframes, njoints, 2)[view_idx] elif FLAGS.mode == 'SMPLMesh': # SMPL Mesh import trimesh # install by `pip install trimesh` import vedo # install by `pip install vedo` smpl_poses, smpl_scaling, smpl_trans = AISTDataset.load_motion( aist_dataset.motion_dir, seq_name) smpl = SMPL(model_path=FLAGS.smpl_dir, gender='MALE', batch_size=1) vertices = smpl.forward( global_orient=torch.from_numpy(smpl_poses[:, 0:1]).float(), body_pose=torch.from_numpy(smpl_poses[:, 1:]).float(), transl=torch.from_numpy(smpl_trans).float(), scaling=torch.from_numpy(smpl_scaling.reshape(1, 1)).float(), ).vertices.detach().numpy()[0] # first frame faces = smpl.faces mesh = trimesh.Trimesh(vertices, faces) mesh.visual.face_colors = [200, 200, 250, 100] keypoints3d = AISTDataset.load_keypoint3d(aist_dataset.keypoint3d_dir, seq_name, use_optim=True) pts = vedo.Points(keypoints3d[0], r=20) # first frame vedo.show(mesh, pts, interactive=True) exit() # Visualize. os.makedirs(FLAGS.save_dir, exist_ok=True) save_path = os.path.join(FLAGS.save_dir, f'{FLAGS.video_name}.mp4') plot_on_video(keypoints2d, video_path, save_path, fps=60)
def main(args): device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') # ========= Define VIBE model ========= # model = VIBE_Demo( seqlen=16, device=device, n_layers=2, hidden_size=1024, add_linear=True, use_residual=True, ).to(device) # ========= Load pretrained weights ========= # pretrained_file = download_ckpt(use_3dpw=False) ckpt = torch.load(pretrained_file, map_location=device) print(f'Performance of pretrained model on 3DPW: {ckpt["performance"]}') ckpt = ckpt['gen_state_dict'] model.load_state_dict(ckpt, strict=False) model.eval() print(f'Loaded pretrained weights from \"{pretrained_file}\"') total_time = time.time() # ========= Run VIBE on crops ========= # print(f'Running VIBE on crops...') vibe_time = time.time() image_folder = args.input_folder dataset = InferenceFromCrops(image_folder=image_folder) orig_height = orig_width = 512 dataloader = DataLoader(dataset, batch_size=args.vibe_batch_size, num_workers=0) with torch.no_grad(): pred_cam, pred_verts, pred_pose, pred_betas, pred_joints3d, norm_joints2d = [], [], [], [], [], [] for batch_num, batch in enumerate(dataloader): print("BATCH:", batch_num) batch = batch.unsqueeze(0) batch = batch.to(device) batch_size, seqlen = batch.shape[:2] output = model(batch)[-1] pred_cam.append(output['theta'][:, :, :3].reshape( batch_size * seqlen, -1)) pred_verts.append(output['verts'].reshape(batch_size * seqlen, -1, 3)) pred_pose.append(output['theta'][:, :, 3:75].reshape( batch_size * seqlen, -1)) pred_betas.append(output['theta'][:, :, 75:].reshape( batch_size * seqlen, -1)) pred_joints3d.append(output['kp_3d'].reshape( batch_size * seqlen, -1, 3)) pred_cam = torch.cat(pred_cam, dim=0) pred_verts = torch.cat(pred_verts, dim=0) pred_pose = torch.cat(pred_pose, dim=0) pred_betas = torch.cat(pred_betas, dim=0) pred_joints3d = torch.cat(pred_joints3d, dim=0) del batch # ========= [Optional] run Temporal SMPLify to refine the results ========= # if args.run_smplify and args.tracking_method == 'pose': norm_joints2d = np.concatenate(norm_joints2d, axis=0) norm_joints2d = convert_kps(norm_joints2d, src='staf', dst='spin') norm_joints2d = torch.from_numpy(norm_joints2d).float().to(device) # Run Temporal SMPLify update, new_opt_vertices, new_opt_cam, new_opt_pose, new_opt_betas, \ new_opt_joints3d, new_opt_joint_loss, opt_joint_loss = smplify_runner( pred_rotmat=pred_pose, pred_betas=pred_betas, pred_cam=pred_cam, j2d=norm_joints2d, device=device, batch_size=norm_joints2d.shape[0], pose2aa=False, ) # update the parameters after refinement print( f'Update ratio after Temporal SMPLify: {update.sum()} / {norm_joints2d.shape[0]}' ) pred_verts = pred_verts.cpu() pred_cam = pred_cam.cpu() pred_pose = pred_pose.cpu() pred_betas = pred_betas.cpu() pred_joints3d = pred_joints3d.cpu() pred_verts[update] = new_opt_vertices[update] pred_cam[update] = new_opt_cam[update] pred_pose[update] = new_opt_pose[update] pred_betas[update] = new_opt_betas[update] pred_joints3d[update] = new_opt_joints3d[update] elif args.run_smplify and args.tracking_method == 'bbox': print( '[WARNING] You need to enable pose tracking to run Temporal SMPLify algorithm!' ) print('[WARNING] Continuing without running Temporal SMPLify!..') # ========= Save results to a pickle file ========= # output_path = image_folder.replace('cropped_frames', 'vibe_results') os.makedirs(output_path, exist_ok=True) pred_cam = pred_cam.cpu().numpy() pred_verts = pred_verts.cpu().numpy() pred_pose = pred_pose.cpu().numpy() pred_betas = pred_betas.cpu().numpy() pred_joints3d = pred_joints3d.cpu().numpy() vibe_results = { 'pred_cam': pred_cam, 'verts': pred_verts, 'pose': pred_pose, 'betas': pred_betas, 'joints3d': pred_joints3d, } del model end = time.time() fps = len(dataset) / (end - vibe_time) print(f'VIBE FPS: {fps:.2f}') total_time = time.time() - total_time print( f'Total time spent: {total_time:.2f} seconds (including model loading time).' ) print( f'Total FPS (including model loading time): {len(dataset) / total_time:.2f}.' ) print( f'Saving vibe results to \"{os.path.join(output_path, "vibe_results.pkl")}\".' ) with open(os.path.join(output_path, "vibe_results.pkl"), 'wb') as f_save: pickle.dump(vibe_results, f_save) if not args.no_render: # ========= Render results as a single video ========= # renderer = Renderer(resolution=(orig_width, orig_height), orig_img=True, wireframe=args.wireframe) output_img_folder = os.path.join(output_path, 'vibe_images') os.makedirs(output_img_folder, exist_ok=True) print(f'Rendering output video, writing frames to {output_img_folder}') image_file_names = sorted([ os.path.join(image_folder, x) for x in os.listdir(image_folder) if x.endswith('.png') or x.endswith('.jpg') ]) for frame_idx in tqdm(range(len(image_file_names))): img_fname = image_file_names[frame_idx] img = cv2.imread(img_fname) frame_verts = vibe_results['verts'][frame_idx] frame_cam = vibe_results['pred_cam'][frame_idx] mesh_filename = None if args.save_obj: mesh_folder = os.path.join(output_path, 'vibe_meshes') os.makedirs(mesh_folder, exist_ok=True) mesh_filename = os.path.join(mesh_folder, f'{frame_idx:06d}.obj') rend_img = renderer.render( img, frame_verts, cam=frame_cam, mesh_filename=mesh_filename, ) whole_img = rend_img if args.sideview: side_img_bg = np.zeros_like(img) side_rend_img90 = renderer.render( side_img_bg, frame_verts, cam=frame_cam, angle=90, axis=[0, 1, 0], ) side_rend_img270 = renderer.render( side_img_bg, frame_verts, cam=frame_cam, angle=270, axis=[0, 1, 0], ) if args.reposed_render: smpl = SMPL('data/vibe_data', batch_size=1) zero_pose = torch.from_numpy( np.zeros((1, pred_pose.shape[-1]))).float() zero_pose[:, 0] = np.pi pred_frame_betas = torch.from_numpy( pred_betas[frame_idx][None, :]).float() with torch.no_grad(): reposed_smpl_output = smpl( betas=pred_frame_betas, body_pose=zero_pose[:, 3:], global_orient=zero_pose[:, :3]) reposed_verts = reposed_smpl_output.vertices reposed_verts = reposed_verts.cpu().detach().numpy() reposed_cam = np.array([0.9, 0, 0]) reposed_rend_img = renderer.render(side_img_bg, reposed_verts[0], cam=reposed_cam) reposed_rend_img90 = renderer.render(side_img_bg, reposed_verts[0], cam=reposed_cam, angle=90, axis=[0, 1, 0]) top_row = np.concatenate( [img, reposed_rend_img, reposed_rend_img90], axis=1) bot_row = np.concatenate( [rend_img, side_rend_img90, side_rend_img270], axis=1) whole_img = np.concatenate([top_row, bot_row], axis=0) else: top_row = np.concatenate([img, side_img_bg, side_img_bg], axis=1) bot_row = np.concatenate( [rend_img, side_rend_img90, side_rend_img270], axis=1) whole_img = np.concatenate([top_row, bot_row], axis=0) # cv2.imwrite(os.path.join(output_img_folder, f'{frame_idx:06d}.png'), whole_img) cv2.imwrite( os.path.join(output_img_folder, os.path.basename(img_fname)), whole_img) # ========= Save rendered video ========= # save_vid_path = os.path.join(output_path, 'vibe_video.mp4') print(f'Saving result video to {save_vid_path}') images_to_video(img_folder=output_img_folder, output_vid_file=save_vid_path) print('================= END =================')