-
Notifications
You must be signed in to change notification settings - Fork 61
/
joints2bvh.py
139 lines (113 loc) · 6.22 KB
/
joints2bvh.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import visualization.Animation as Animation
from visualization.InverseKinematics import BasicInverseKinematics, BasicJacobianIK, InverseKinematics
from visualization.Quaternions import Quaternions
import visualization.BVH_mod as BVH
from visualization.remove_fs import *
from utils.plot_script import plot_3d_motion
from utils import paramUtil
from common.skeleton import Skeleton
import torch
from torch import nn
from visualization.utils.quat import ik_rot, between, fk, ik
from tqdm import tqdm
def get_grot(glb, parent, offset):
root_quat = np.array([[1.0, 0.0, 0.0, 0.0]]).repeat(glb.shape[0], axis=0)[:, None]
local_pos = glb[:, 1:] - glb[:, parent[1:]]
norm_offset = offset[1:] / np.linalg.norm(offset[1:], axis=-1, keepdims=True)
norm_lpos = local_pos / np.linalg.norm(local_pos, axis=-1, keepdims=True)
grot = between(norm_offset, norm_lpos)
grot = np.concatenate((root_quat, grot), axis=1)
grot /= np.linalg.norm(grot, axis=-1, keepdims=True)
return grot
class Joint2BVHConvertor:
def __init__(self):
self.template = BVH.load('./visualization/data/template.bvh', need_quater=True)
self.re_order = [0, 1, 4, 7, 10, 2, 5, 8, 11, 3, 6, 9, 12, 15, 13, 16, 18, 20, 14, 17, 19, 21]
self.re_order_inv = [0, 1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8, 12, 14, 18, 13, 15, 19, 16, 20, 17, 21]
self.end_points = [4, 8, 13, 17, 21]
self.template_offset = self.template.offsets.copy()
self.parents = [-1, 0, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 12, 11, 14, 15, 16, 11, 18, 19, 20]
def convert(self, positions, filename, iterations=10, foot_ik=True):
'''
Convert the SMPL joint positions to Mocap BVH
:param positions: (N, 22, 3)
:param filename: Save path for resulting BVH
:param iterations: iterations for optimizing rotations, 10 is usually enough
:param foot_ik: whether to enfore foot inverse kinematics, removing foot slide issue.
:return:
'''
positions = positions[:, self.re_order]
new_anim = self.template.copy()
new_anim.rotations = Quaternions.id(positions.shape[:-1])
new_anim.positions = new_anim.positions[0:1].repeat(positions.shape[0], axis=-0)
new_anim.positions[:, 0] = positions[:, 0]
if foot_ik:
positions = remove_fs(positions, None, fid_l=(3, 4), fid_r=(7, 8), interp_length=5,
force_on_floor=True)
ik_solver = BasicInverseKinematics(new_anim, positions, iterations=iterations, silent=True)
new_anim = ik_solver()
# BVH.save(filename, new_anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True)
glb = Animation.positions_global(new_anim)[:, self.re_order_inv]
if filename is not None:
BVH.save(filename, new_anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True)
return new_anim, glb
def convert_sgd(self, positions, filename, iterations=100, foot_ik=True):
'''
Convert the SMPL joint positions to Mocap BVH
:param positions: (N, 22, 3)
:param filename: Save path for resulting BVH
:param iterations: iterations for optimizing rotations, 10 is usually enough
:param foot_ik: whether to enfore foot inverse kinematics, removing foot slide issue.
:return:
'''
## Positional Foot locking ##
glb = positions[:, self.re_order]
if foot_ik:
glb = remove_fs(glb, None, fid_l=(3, 4), fid_r=(7, 8), interp_length=2,
force_on_floor=True)
## Fit BVH ##
new_anim = self.template.copy()
new_anim.rotations = Quaternions.id(glb.shape[:-1])
new_anim.positions = new_anim.positions[0:1].repeat(glb.shape[0], axis=-0)
new_anim.positions[:, 0] = glb[:, 0]
anim = new_anim.copy()
rot = torch.tensor(anim.rotations.qs, dtype=torch.float)
pos = torch.tensor(anim.positions[:, 0, :], dtype=torch.float)
offset = torch.tensor(anim.offsets, dtype=torch.float)
glb = torch.tensor(glb, dtype=torch.float)
ik_solver = InverseKinematics(rot, pos, offset, anim.parents, glb)
print('Fixing foot contact using IK...')
for i in tqdm(range(iterations)):
mse = ik_solver.step()
# print(i, mse)
rotations = ik_solver.rotations.detach().cpu()
norm = torch.norm(rotations, dim=-1, keepdim=True)
rotations /= norm
anim.rotations = Quaternions(rotations.numpy())
anim.rotations[:, self.end_points] = Quaternions.id((anim.rotations.shape[0], len(self.end_points)))
anim.positions[:, 0, :] = ik_solver.position.detach().cpu().numpy()
if filename is not None:
BVH.save(filename, anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True)
# BVH.save(filename[:-3] + 'bvh', anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True)
glb = Animation.positions_global(anim)[:, self.re_order_inv]
return anim, glb
if __name__ == "__main__":
# file = 'batch0_sample13_repeat0_len196.npy'
# file = 'batch2_sample10_repeat0_len156.npy'
# file = 'batch2_sample13_repeat0_len196.npy' #line #57 new_anim.positions = lpos #new_anim.positions[0:1].repeat(positions.shape[0], axis=-0) #TODO, figure out why it's important
# file = 'batch1_sample12_repeat0_len196.npy' #hard case karate
# file = 'batch1_sample14_repeat0_len180.npy'
# file = 'batch0_sample3_repeat0_len192.npy'
# file = 'batch1_sample4_repeat0_len136.npy'
# file = 'batch0_sample0_repeat0_len152.npy'
# path = f'/Users/yuxuanmu/project/MaskMIT/demo/cond4_topkr0.9_ts18_tau1.0_s1009/joints/{file}'
# joints = np.load(path)
# converter = Joint2BVHConvertor()
# new_anim = converter.convert(joints, './gen_L196.mp4', foot_ik=True)
folder = '/Users/yuxuanmu/project/MaskMIT/demo/cond4_topkr0.9_ts18_tau1.0_s1009'
files = os.listdir(os.path.join(folder, 'joints'))
files = [f for f in files if 'repeat' in f]
converter = Joint2BVHConvertor()
for f in tqdm(files):
joints = np.load(os.path.join(folder, 'joints', f))
converter.convert(joints, os.path.join(folder, 'ik_animations', f'ik_{f}'.replace('npy', 'mp4')), foot_ik=True)