Beispiel #1
0
 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
Beispiel #2
0
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)