def smpl_to_kpts(self, pred_rotmat, pred_shape, pred_cam, J_regressor): pred_output = self.smpl(betas=pred_shape, body_pose=pred_rotmat[:, 1:], global_orient=pred_rotmat[:, 0].unsqueeze(1), pose2rot=False) pred_vertices = pred_output.vertices pred_joints = pred_output.joints if J_regressor is not None: J_regressor_batch = J_regressor[None, :].expand( pred_vertices.shape[0], -1, -1).to(pred_vertices.device) pred_joints = torch.matmul(J_regressor_batch, pred_vertices) pred_joints = pred_joints[:, H36M_TO_J14, :] pred_keypoints_2d = projection(pred_joints, pred_cam) pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape( -1, 72) output = [{ 'theta': torch.cat([pred_cam, pose, pred_shape], dim=1), 'verts': pred_vertices, 'kp_2d': pred_keypoints_2d, 'kp_3d': pred_joints, 'rotmat': pred_rotmat }] return output
def forward(self, x, init_pose=None, init_shape=None, init_cam=None, n_iter=3, J_regressor=None): batch_size = x.shape[0] if init_pose is None: init_pose = self.init_pose.expand(batch_size, -1) if init_shape is None: init_shape = self.init_shape.expand(batch_size, -1) if init_cam is None: init_cam = self.init_cam.expand(batch_size, -1) pred_pose = init_pose pred_shape = init_shape pred_cam = init_cam for i in range(n_iter): xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1) xc = self.fc1(xc) xc = self.drop1(xc) xc = self.fc2(xc) xc = self.drop2(xc) pred_pose = self.decpose(xc) + pred_pose pred_shape = self.decshape(xc) + pred_shape pred_cam = self.deccam(xc) + pred_cam pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) pred_output = self.smpl(betas=pred_shape, body_pose=pred_rotmat[:, 1:], global_orient=pred_rotmat[:, 0].unsqueeze(1), pose2rot=False) pred_vertices = pred_output.vertices pred_joints = pred_output.joints if J_regressor is not None: J_regressor_batch = J_regressor[None, :].expand( pred_vertices.shape[0], -1, -1).to(pred_vertices.device) pred_joints = torch.matmul(J_regressor_batch, pred_vertices) pred_joints = pred_joints[:, H36M_TO_J14, :] pred_keypoints_2d = projection(pred_joints, pred_cam) pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape( -1, 72) output = [{ 'theta': torch.cat([pred_cam, pose, pred_shape], dim=1), 'verts': pred_vertices, 'kp_2d': pred_keypoints_2d, 'kp_3d': pred_joints, 'rotmat': pred_rotmat }] return output
def forward(self, input, J_regressor=None): batch_size, seqlen, nc, h, w = input.shape feature = self.hmr.feature_extractor(input.reshape(-1, nc, h, w)) input = feature.reshape(batch_size, seqlen, -1) # input size NTF batch_size, seqlen = input.shape[:2] vibe_output = self.vibe(input)[-1] vibe_theta = vibe_output['theta'] vibe_pose = vibe_theta[:,:,3:75] pred_cam = vibe_theta[:,:,:3].reshape(batch_size * seqlen, 3) pred_shape = vibe_theta[:,:,75:].reshape(batch_size * seqlen, 10) vibe_pose_6d = convert_aa_to_orth6d(vibe_pose).reshape(vibe_pose.shape[0], vibe_pose.shape[1], -1) X_r= self.refiner_model(vibe_pose_6d.permute(1, 0, 2), input.permute(1, 0, 2)) X_r = X_r.permute(1, 0, 2).contiguous()[:,:seqlen, :] pred_rotmat = convert_orth_6d_to_mat(X_r).reshape(batch_size * seqlen, 24, 3, 3) # pred_rotmat = convert_orth_6d_to_mat(vibe_pose_6d).reshape(batch_size * seqlen, 24, 3, 3) pred_output = self.smpl( betas=pred_shape, body_pose=pred_rotmat[:, 1:], global_orient=pred_rotmat[:, 0:1], pose2rot=False ) pred_vertices = pred_output.vertices pred_joints = pred_output.joints if J_regressor is not None: J_regressor_batch = J_regressor[None, :].expand(pred_vertices.shape[0], -1, -1).to(pred_vertices.device) pred_joints = torch.matmul(J_regressor_batch, pred_vertices) pred_joints = pred_joints[:, H36M_TO_J14, :] pred_keypoints_2d = projection(pred_joints, pred_cam) pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape(-1, 72) smpl_output = [{ 'theta' : torch.cat([pred_cam, pose, pred_shape], dim=1), 'verts' : pred_vertices, 'kp_2d' : pred_keypoints_2d, 'kp_3d' : pred_joints, 'rotmat' : pred_rotmat }] for s in smpl_output: s['theta'] = s['theta'].reshape(batch_size, seqlen, -1) s['verts'] = s['verts'].reshape(batch_size, seqlen, -1, 3) s['kp_2d'] = s['kp_2d'].reshape(batch_size, seqlen, -1, 2) s['kp_3d'] = s['kp_3d'].reshape(batch_size, seqlen, -1, 3) s['rotmat'] = s['rotmat'].reshape(batch_size, seqlen, -1, 3, 3) return smpl_output
def forward(self, input, J_regressor=None): # input size NTF batch_size, seqlen, nc, h, w = input.shape feature = self.hmr.feature_extractor(input.reshape(-1, nc, h, w)) input = feature.reshape(batch_size, seqlen, -1) batch_size, seqlen = input.shape[:2] X_r, pred_cam, pred_shape = self.meva_model(input.permute(1, 0, 2)) X_r, pred_cam, pred_shape = X_r.permute(1, 0, 2)[:,:seqlen,:].reshape(batch_size * seqlen, -1), \ pred_cam.permute(1, 0, 2)[:,:seqlen,:].reshape(batch_size * seqlen, -1), \ pred_shape.permute(1, 0, 2)[:,:seqlen,:].reshape(batch_size * seqlen, -1) # pred_shape = pred_shape[:, None, :].expand(batch_size, seqlen, -1).reshape(batch_size * seqlen, -1) pred_rotmat = convert_orth_6d_to_mat(X_r).view(batch_size * seqlen, 24, 3, 3) pred_output = self.smpl(betas=pred_shape, body_pose=pred_rotmat[:, 1:], global_orient=pred_rotmat[:, 0].unsqueeze(1), pose2rot=False) pred_vertices = pred_output.vertices pred_joints = pred_output.joints if J_regressor is not None: J_regressor_batch = J_regressor[None, :].expand( pred_vertices.shape[0], -1, -1).to(pred_vertices.device) pred_joints = torch.matmul(J_regressor_batch, pred_vertices) pred_joints = pred_joints[:, H36M_TO_J14, :] pred_keypoints_2d = projection(pred_joints, pred_cam) pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape( -1, 72) smpl_output = [{ 'theta': torch.cat([pred_cam, pose, pred_shape], dim=1), 'verts': pred_vertices, 'kp_2d': pred_keypoints_2d, 'kp_3d': pred_joints, 'rotmat': pred_rotmat }] for s in smpl_output: s['theta'] = s['theta'].reshape(batch_size, seqlen, -1) s['verts'] = s['verts'].reshape(batch_size, seqlen, -1, 3) s['kp_2d'] = s['kp_2d'].reshape(batch_size, seqlen, -1, 2) s['kp_3d'] = s['kp_3d'].reshape(batch_size, seqlen, -1, 3) s['rotmat'] = s['rotmat'].reshape(batch_size, seqlen, -1, 3, 3) return smpl_output
def read_data(folder, set, debug=False): dataset = { 'vid_name': [], 'frame_id': [], 'joints3D': [], 'joints2D': [], 'shape': [], 'pose': [], 'bbox': [], 'img_name': [], 'features': [], 'valid': [], } model = spin.get_pretrained_hmr() if set == 'val': set = 'test' sequences = [ x.split('.')[0] for x in os.listdir(osp.join(folder, 'sequenceFiles', set)) ] J_regressor = None smpl = SMPL(SMPL_MODEL_DIR, batch_size=1, create_transl=False) if set == 'test': J_regressor = torch.from_numpy( np.load(osp.join(VIBE_DATA_DIR, 'J_regressor_h36m.npy'))).float() for i, seq in tqdm(enumerate(sequences)): data_file = osp.join(folder, 'sequenceFiles', set, seq + '.pkl') data = pkl.load(open(data_file, 'rb'), encoding='latin1') img_dir = osp.join(folder, 'imageFiles', seq) num_people = len(data['poses']) num_frames = len(data['img_frame_ids']) assert (data['poses2d'][0].shape[0] == num_frames) for p_id in range(num_people): pose = torch.from_numpy(data['poses'][p_id]).float() shape = torch.from_numpy(data['betas'][p_id][:10]).float().repeat( pose.size(0), 1) trans = torch.from_numpy(data['trans'][p_id]).float() j2d = data['poses2d'][p_id].transpose(0, 2, 1) cam_pose = data['cam_poses'] campose_valid = data['campose_valid'][p_id] # ======== Align the mesh params ======== # rot = pose[:, :3] rot_mat = batch_rodrigues(rot) Rc = torch.from_numpy(cam_pose[:, :3, :3]).float() Rs = torch.bmm(Rc, rot_mat.reshape(-1, 3, 3)) rot = rotation_matrix_to_angle_axis(Rs) pose[:, :3] = rot # ======== Align the mesh params ======== # output = smpl(betas=shape, body_pose=pose[:, 3:], global_orient=pose[:, :3], transl=trans) # verts = output.vertices j3d = output.joints if J_regressor is not None: vertices = output.vertices J_regressor_batch = J_regressor[None, :].expand( vertices.shape[0], -1, -1).to(vertices.device) j3d = torch.matmul(J_regressor_batch, vertices) j3d = j3d[:, H36M_TO_J14, :] img_paths = [] for i_frame in range(num_frames): img_path = os.path.join(img_dir + '/image_{:05d}.jpg'.format(i_frame)) img_paths.append(img_path) bbox_params, time_pt1, time_pt2 = get_smooth_bbox_params( j2d, vis_thresh=VIS_THRESH, sigma=8) # process bbox_params c_x = bbox_params[:, 0] c_y = bbox_params[:, 1] scale = bbox_params[:, 2] w = h = 150. / scale w = h = h * 1.1 bbox = np.vstack([c_x, c_y, w, h]).T # process keypoints j2d[:, :, 2] = j2d[:, :, 2] > 0.3 # set the visibility flags # Convert to common 2d keypoint format perm_idxs = get_perm_idxs('3dpw', 'common') perm_idxs += [0, 0] # no neck, top head j2d = j2d[:, perm_idxs] j2d[:, 12:, 2] = 0.0 # print('j2d', j2d[time_pt1:time_pt2].shape) # print('campose', campose_valid[time_pt1:time_pt2].shape) img_paths_array = np.array(img_paths)[time_pt1:time_pt2] dataset['vid_name'].append( np.array([f'{seq}_{p_id}'] * num_frames)[time_pt1:time_pt2]) dataset['frame_id'].append( np.arange(0, num_frames)[time_pt1:time_pt2]) dataset['img_name'].append(img_paths_array) dataset['joints3D'].append(j3d.numpy()[time_pt1:time_pt2]) dataset['joints2D'].append(j2d[time_pt1:time_pt2]) dataset['shape'].append(shape.numpy()[time_pt1:time_pt2]) dataset['pose'].append(pose.numpy()[time_pt1:time_pt2]) dataset['bbox'].append(bbox) dataset['valid'].append(campose_valid[time_pt1:time_pt2]) features = extract_features(model, img_paths_array, bbox, kp_2d=j2d[time_pt1:time_pt2], debug=debug, dataset='3dpw', scale=1.2) dataset['features'].append(features) for k in dataset.keys(): dataset[k] = np.concatenate(dataset[k]) print(k, dataset[k].shape) # Filter out keypoints indices_to_use = np.where( (dataset['joints2D'][:, :, 2] > VIS_THRESH).sum(-1) > MIN_KP)[0] for k in dataset.keys(): dataset[k] = dataset[k][indices_to_use] return dataset
def smplify_runner(pred_rotmat, pred_betas, pred_cam, j2d, device, batch_size, lr=1.0, opt_steps=1, use_lbfgs=True, pose2aa=True): smplify = TemporalSMPLify( step_size=lr, batch_size=batch_size, num_iters=opt_steps, focal_length=5000., use_lbfgs=use_lbfgs, device=device, # max_iter=10, ) # Convert predicted rotation matrices to axis-angle if pose2aa: pred_pose = rotation_matrix_to_angle_axis( pred_rotmat.detach()).reshape(batch_size, -1) else: pred_pose = pred_rotmat # Calculate camera parameters for smplify pred_cam_t = torch.stack([ pred_cam[:, 1], pred_cam[:, 2], 2 * 5000 / (224 * pred_cam[:, 0] + 1e-9) ], dim=-1) gt_keypoints_2d_orig = j2d # Before running compute reprojection error of the network opt_joint_loss = smplify.get_fitting_loss( pred_pose.detach(), pred_betas.detach(), pred_cam_t.detach(), 0.5 * 224 * torch.ones(batch_size, 2, device=device), gt_keypoints_2d_orig).mean(dim=-1) best_prediction_id = torch.argmin(opt_joint_loss).item() pred_betas = pred_betas[best_prediction_id].unsqueeze(0) # pred_betas = pred_betas[best_prediction_id:best_prediction_id+2] # .unsqueeze(0) # top5_best_idxs = torch.topk(opt_joint_loss, 5, largest=False)[1] # breakpoint() start = time.time() # Run SMPLify optimization initialized from the network prediction # new_opt_vertices, new_opt_joints, \ # new_opt_pose, new_opt_betas, \ # new_opt_cam_t, \ output, new_opt_joint_loss = smplify( pred_pose.detach(), pred_betas.detach(), pred_cam_t.detach(), 0.5 * 224 * torch.ones(batch_size, 2, device=device), gt_keypoints_2d_orig, ) new_opt_joint_loss = new_opt_joint_loss.mean(dim=-1) # smplify_time = time.time() - start # print(f'Smplify time: {smplify_time}') # Will update the dictionary for the examples where the new loss is less than the current one update = (new_opt_joint_loss < opt_joint_loss) new_opt_vertices = output['verts'] new_opt_cam_t = output['theta'][:, :3] new_opt_pose = output['theta'][:, 3:75] new_opt_betas = output['theta'][:, 75:] return_val = [ update, new_opt_vertices.cpu(), new_opt_cam_t.cpu(), new_opt_pose.cpu(), new_opt_betas.cpu(), new_opt_joint_loss, opt_joint_loss, ] return return_val
def forward(self, x, init_pose=None, init_shape=None, init_cam=None, n_iter=3, return_features=False): batch_size = x.shape[0] if init_pose is None: init_pose = self.init_pose.expand(batch_size, -1) if init_shape is None: init_shape = self.init_shape.expand(batch_size, -1) if init_cam is None: init_cam = self.init_cam.expand(batch_size, -1) x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x1 = self.layer1(x) x2 = self.layer2(x1) x3 = self.layer3(x2) x4 = self.layer4(x3) xf = self.avgpool(x4) xf = xf.view(xf.size(0), -1) pred_pose = init_pose pred_shape = init_shape pred_cam = init_cam for i in range(n_iter): xc = torch.cat([xf, pred_pose, pred_shape, pred_cam], 1) xc = self.fc1(xc) xc = self.drop1(xc) xc = self.fc2(xc) xc = self.drop2(xc) pred_pose = self.decpose(xc) + pred_pose pred_shape = self.decshape(xc) + pred_shape pred_cam = self.deccam(xc) + pred_cam pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) pred_output = self.smpl(betas=pred_shape, body_pose=pred_rotmat[:, 1:], global_orient=pred_rotmat[:, 0].unsqueeze(1), pose2rot=False) pred_vertices = pred_output.vertices pred_joints = pred_output.joints pred_keypoints_2d = projection(pred_joints, pred_cam) pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape( -1, 72) output = [{ 'theta': torch.cat([pred_cam, pose, pred_shape], dim=1), 'verts': pred_vertices, 'kp_2d': pred_keypoints_2d, 'kp_3d': pred_joints, }] if return_features: return xf, output else: return output
def forward(self, x, init_pose=None, init_shape=None, init_cam=None, n_iter=3, J_regressor=None): batch_size = x.shape[0] if init_pose is None: init_pose = self.init_pose.expand(batch_size, -1) if init_shape is None: init_shape = self.init_shape.expand(batch_size, -1) if init_cam is None: init_cam = self.init_cam.expand(batch_size, -1) pred_pose = init_pose pred_shape = init_shape pred_cam = init_cam for i in range(n_iter): xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1) xc = self.fc1(xc) xc = self.drop1(xc) xc = self.fc2(xc) xc = self.drop2(xc) pred_pose = self.decpose(xc) + pred_pose pred_shape = self.decshape(xc) + pred_shape pred_cam = self.deccam(xc) + pred_cam # print('inside SPIN model, pred pose shape',pred_pose.shape) pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) # print("Inputto spin model pose--",pred_pose.shape) # print("Inputto spin model ---- ",pred_shape.shape,pred_shape,pred_rotmat[:, 1:].shape,pred_rotmat[:, 0].unsqueeze(1).shape) pred_output = self.smpl(betas=pred_shape, body_pose=pred_rotmat[:, 1:], global_orient=pred_rotmat[:, 0].unsqueeze(1), pose2rot=False) pred_vertices = pred_output.vertices pred_joints = pred_output.joints #print('inside spin model,pred_vertices shape ',pred_joints.shape) #H36M_TO_J17 = [8,5,45,46,4,7,21,19,17,16,18,20,47,48,51,50,24] #H36M_TO_J17 = [3,2,1,6,7,8,27,26,25,17,18,19,14,15,12,13,14] #H36M_TO_J17 = [8,5,2,1,4,7,21,19,17,16,18,20,12,15,3,9,15] H36M_TO_J17 = [ 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 41, 40, 44 ] #print("@@@@@ USED!!!!") #print('predicted joints shape',pred_joints.shape) pred_joints = pred_joints[:, H36M_TO_J17, :] #print('after predicted joints shape',pred_joints.shape) if J_regressor is not None: J_regressor_batch = J_regressor[None, :].expand( pred_vertices.shape[0], -1, -1).to(pred_vertices.device) pred_joints = torch.matmul(J_regressor_batch, pred_vertices) pred_joints = pred_joints[:, H36M_TO_J14, :] #print('after predicted joints shape',pred_joints.shape) pred_keypoints_2d = projection(pred_joints, pred_cam) pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape( -1, 72) output = [{ 'theta': torch.cat([pred_cam, pose, pred_shape], dim=1), 'verts': pred_vertices, 'kp_2d': pred_keypoints_2d, 'kp_3d': pred_joints, 'rotmat': pred_rotmat }] return output