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)
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)