def _valid_epoch(self, epoch): self.model.eval() total_val_metrics = 0 total_val_loss = 0 for batch_idx, datas in enumerate(self.test_data_loader): datas = [ self._prepare_data(item, _from='tensor') for item in datas[:-1] ] poses_2d_pixel, poses_2d, poses_3d, bones, contacts, alphas, proj_facters = datas _, _, _, fake_pose_3d, _, _ = self.model.forward_fk( poses_2d, self.test_parameters) total_val_metrics += metric.mean_points_error( fake_pose_3d, poses_3d) * torch.mean( alphas[0]).data.cpu().numpy() total_val_loss += torch.mean( torch.norm(fake_pose_3d.view(-1, 17, 3) - poses_3d.view(-1, 17, 3), dim=-1)).item() val_log = { 'val_metric': total_val_metrics / len(self.test_data_loader), 'val_loss': total_val_loss / len(self.test_data_loader), } self.writer.set_step(epoch, mode='valid') self.writer.set_scalars(val_log) self.train_logger.info( 'Eveluation: mean_points_error: {:.6f} loss: {:.6f}'.format( val_log['val_metric'], val_log['val_loss'])) return val_log
def main(config, args, output_folder): resume = args.resume name_list = [ 'Hips', 'RightUpLeg', 'RightLeg', 'RightFoot', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightArm', 'RightForeArm', 'RightHand' ] model = getattr(models, config.arch.type)(config) checkpoint = torch.load(resume) state_dict = checkpoint['state_dict'] model.load_state_dict(state_dict) device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') model = model.to(device) model.eval() if args.input == 'h36m': test_data_loader = h36m_loader(config, is_training=False) test_parameters = [ torch.from_numpy(np.array(item)).float().to(device) for item in test_data_loader.dataset.get_parameters() ] error_list = {} errors = [] sampling_export = np.random.choice(test_data_loader.n_samples - 1, size=4, replace=False) for video_idx, datas in enumerate(test_data_loader): video_name = datas[-1][0] datas = [item.float().to(device) for item in datas[:-1]] poses_2d, poses_3d, bones, contacts, alphas, proj_facters = datas with torch.no_grad(): pre_bones, pre_rotations, pre_rotations_full, pre_pose_3d, pre_c, pre_proj = model.forward_fk( poses_2d, test_parameters) error = metric.mean_points_error(poses_3d, pre_pose_3d) * torch.mean( alphas[0]).data.cpu().numpy() errors.append(error) action_name = video_name.split('_')[1].split(' ')[0] if action_name in error_list.keys(): error_list[action_name].append(error) else: error_list[action_name] = [error] if video_idx in sampling_export: if config.arch.translation: R, T, f, c, k, p, res_w, res_h = test_data_loader.dataset.cameras[ (int(video_name.split('_')[0].replace('S', '')), int(video_name.split('_')[-1]))] pose_2d_film = (poses_2d[0, :, :2].cpu().numpy() - c[:, 0]) / f[:, 0] translations = np.ones(shape=(pose_2d_film.shape[0], 3)) translations[:, :2] = pose_2d_film translation = (translations * np.repeat( pre_proj[0].cpu().numpy(), 3, axis=-1).reshape( (-1, 3))) * 5 else: translation = np.zeros((poses_2d.shape[1], 3)) rotations = pre_rotations_full[0].cpu().numpy() length = (pre_bones * test_parameters[3].unsqueeze(0) + test_parameters[2].repeat(bones.shape[0], 1, 1))[0].cpu().numpy() BVH.save('%s/%s.bvh' % (output_folder, video_name), Animation.load_from_network(translation, rotations, length, third_dimension=1), names=name_list) error_file = '%s/errors.txt' % output_folder with open(error_file, 'w') as f: f.writelines('=====Action===== ==mm==\n') total = [] for key in error_list.keys(): mean_error = np.mean(np.array(error_list[key])) total.append(mean_error) print('%16s %.2f' % (key, mean_error)) f.writelines('%16s %.2f \n' % (key, mean_error)) print('%16s %.2f' % ('Average', np.mean(np.array(errors)))) f.writelines('%16s %.2f \n' % ('Average', np.mean(np.array(errors)))) f.close() else: parameters = [ torch.from_numpy(np.array(item)).float().to(device) for item in h36m_loader(config, is_training=True).dataset.get_parameters() ] def export(pose_folder): video_name = pose_folder.split('/')[-1] files = util.make_dataset([pose_folder], phase='json', data_split=1, sort=True, sort_index=0) IMAGE_WIDTH = 1080 # Should be changed refer to your test data pose_batch = [] confidence_batch = [] for pose_file_name in files: with open(pose_file_name, 'r') as f: h36m_locations, h36m_confidence = h36m_utils.convert_openpose( json.load(f)) pose_batch.append(h36m_locations) confidence_batch.append(h36m_confidence) poses_2d = np.concatenate(pose_batch, axis=0) / IMAGE_WIDTH confidences = np.concatenate(confidence_batch, axis=0) poses_2d_root = ( poses_2d - np.tile(poses_2d[:, :2], [1, int(poses_2d.shape[-1] / 2)])) if config.arch.confidence: poses_2d_root_c = np.zeros( (poses_2d_root.shape[0], int(poses_2d_root.shape[-1] / 2 * 3))) for joint_index in range(int(poses_2d_root.shape[-1] / 2)): poses_2d_root_c[:, 3 * joint_index] = poses_2d_root[:, 2 * joint_index].copy( ) poses_2d_root_c[:, 3 * joint_index + 1] = poses_2d_root[:, 2 * joint_index + 1].copy() poses_2d_root_c[:, 3 * joint_index + 2] = np.array( confidences)[:, joint_index].copy() poses_2d = poses_2d_root_c poses_2d = np.divide((poses_2d - parameters[0].cpu().numpy()), parameters[1].cpu().numpy()) poses_2d = torch.from_numpy( np.array(poses_2d)).unsqueeze(0).float().to(device) with torch.no_grad(): pre_bones, pre_rotations, pre_rotations_full, pre_pose_3d, pre_c, pre_proj = model.forward_fk( poses_2d, parameters) if config.arch.translation: pose_2d_film = (poses_2d[0, :, :2].cpu().numpy() - 0.5) translations = np.ones(shape=(pose_2d_film.shape[0], 3)) translations[:, :2] = pose_2d_film translation = (translations * np.repeat( pre_proj[0].cpu().numpy(), 3, axis=-1).reshape( (-1, 3))) * 3 translation[:] -= translation[[0]] else: translation = np.zeros((poses_2d.shape[1], 3)) rotations = pre_rotations_full[0].cpu().numpy() length = (pre_bones * parameters[3].unsqueeze(0) + parameters[2].repeat(pre_bones.shape[0], 1, 1))[0].cpu().numpy() BVH.save('%s/%s.bvh' % (output_folder, video_name), Animation.load_from_network(translation, rotations, length, third_dimension=1), names=name_list) print('The bvh file of %s has been saved!' % video_name) export(args.input)