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 remove_fs(anim,
foot,
output_path,
fid_l=(4, 5),
fid_r=(9, 10),
interp_length=5,
force_on_floor=True):
(anim, names, ftime), glb = nrot2anim(anim)
T = len(glb)
fid = list(fid_l) + list(fid_r)
fid_l, fid_r = np.array(fid_l), np.array(fid_r)
foot_heights = np.minimum(glb[:, fid_l, 1],
glb[:, fid_r, 1]).min(axis=1) # [T, 2] -> [T]
# print(np.min(foot_heights))
floor_height = softmin(foot_heights, softness=0.5, axis=0)
# print(floor_height)
glb[:, :, 1] -= floor_height
anim.positions[:, 0, 1] -= floor_height
for i, fidx in enumerate(fid):
fixed = foot[i] # [T]
"""
for t in range(T):
glb[t, fidx][1] = max(glb[t, fidx][1], 0.25)
"""
s = 0
while s < T:
while s < T and fixed[s] == 0:
s += 1
if s >= T:
break
t = s
avg = glb[t, fidx].copy()
while t + 1 < T and fixed[t + 1] == 1:
t += 1
avg += glb[t, fidx].copy()
avg /= (t - s + 1)
if force_on_floor:
avg[1] = 0.0
for j in range(s, t + 1):
glb[j, fidx] = avg.copy()
# print(fixed[s - 1:t + 2])
s = t + 1
for s in range(T):
if fixed[s] == 1:
continue
l, r = None, None
consl, consr = False, False
for k in range(interp_length):
if s - k - 1 < 0:
break
if fixed[s - k - 1]:
l = s - k - 1
consl = True
break
for k in range(interp_length):
if s + k + 1 >= T:
break
if fixed[s + k + 1]:
r = s + k + 1
consr = True
break
if not consl and not consr:
continue
if consl and consr:
litp = lerp(alpha(1.0 * (s - l + 1) / (interp_length + 1)),
glb[s, fidx], glb[l, fidx])
ritp = lerp(alpha(1.0 * (r - s + 1) / (interp_length + 1)),
glb[s, fidx], glb[r, fidx])
itp = lerp(alpha(1.0 * (s - l + 1) / (r - l + 1)), ritp, litp)
glb[s, fidx] = itp.copy()
continue
if consl:
litp = lerp(alpha(1.0 * (s - l + 1) / (interp_length + 1)),
glb[s, fidx], glb[l, fidx])
glb[s, fidx] = litp.copy()
continue
if consr:
ritp = lerp(alpha(1.0 * (r - s + 1) / (interp_length + 1)),
glb[s, fidx], glb[r, fidx])
glb[s, fidx] = ritp.copy()
targetmap = {}
for j in range(glb.shape[1]):
targetmap[j] = glb[:, j]
ik = JacobianInverseKinematics(anim,
targetmap,
iterations=10,
damping=4.0,
silent=False)
ik()
if not os.path.exists(os.path.dirname(output_path)):
os.makedirs(os.path.dirname(output_path))
BVH.save(output_path, anim, names, ftime)
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)
def save_bvh_from_network_output(nrot, output_path):
anim = AnimationData.from_network_output(nrot)
bvh, names, ftime = anim.get_BVH()
if not os.path.exists(os.path.dirname(output_path)):
os.makedirs(os.path.dirname(output_path))
BVH.save(output_path, bvh, names, ftime)
