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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224 | import torch
import bvh.bvh_io as bvh_io
import numpy as np
from bvh.Quaternions import Quaternions
from bvh.skeleton_database import SkeletonDatabase
from models.kinematics import ForwardKinematicsJoint
from models.transforms import quat2repr6d, quat2mat
from models.contact import foot_contact
from bvh.bvh_writer import WriterWrapper
class Skeleton:
def __init__(self, names, parent, offsets, joint_reduction=True):
self._names = names
self.original_parent = parent
self._offsets = offsets
self._parent = None
self._ee_id = None
self.contact_names = []
for i, name in enumerate(self._names):
if ':' in name:
self._names[i] = name[name.find(':')+1:]
if joint_reduction:
self.skeleton_type, match_num = SkeletonDatabase.match(names)
corps_names = SkeletonDatabase.corps_names[self.skeleton_type]
self.contact_names = SkeletonDatabase.contact_names[self.skeleton_type]
self.contact_threshold = SkeletonDatabase.contact_thresholds[self.skeleton_type]
self.contact_id = []
for i in self.contact_names:
self.contact_id.append(corps_names.index(i))
else:
self.skeleton_type = -1
corps_names = self._names
self.details = [] # joints that does not belong to the corps (we are not interested in them)
for i, name in enumerate(self._names):
if name not in corps_names: self.details.append(i)
self.corps = []
self.simplified_name = []
self.simplify_map = {}
self.inverse_simplify_map = {}
# Repermute the skeleton id according to the databse
for name in corps_names:
for j in range(len(self._names)):
if name in self._names[j]:
self.corps.append(j)
break
if len(self.corps) != len(corps_names):
for i in self.corps:
print(self._names[i], end=' ')
print(self.corps, self.skeleton_type, len(self.corps), sep='\n')
raise Exception('Problem in this skeleton')
self.joint_num_simplify = len(self.corps)
for i, j in enumerate(self.corps):
self.simplify_map[j] = i
self.inverse_simplify_map[i] = j
self.simplified_name.append(self._names[j])
self.inverse_simplify_map[0] = -1
for i in range(len(self._names)):
if i in self.details:
self.simplify_map[i] = -1
@property
def parent(self):
if self._parent is None:
self._parent = self.original_parent[self.corps].copy()
for i in range(self._parent.shape[0]):
if i >= 1: self._parent[i] = self.simplify_map[self._parent[i]]
self._parent = tuple(self._parent)
return self._parent
@property
def offsets(self):
return torch.tensor(self._offsets[self.corps], dtype=torch.float)
@property
def names(self):
return self.simplified_name
@property
def ee_id(self):
raise Exception('Abaddoned')
# if self._ee_id is None:
# self._ee_id = []
# for i in SkeletonDatabase.ee_names[self.skeleton_type]:
# self.ee_id._ee_id(corps_names[self.skeleton_type].index(i))
class BVH_file:
def __init__(self, file_path, no_scale=False, requires_contact=False, joint_reduction=True):
self.anim = bvh_io.load(file_path)
self._names = self.anim.names
self.frametime = self.anim.frametime
self.skeleton = Skeleton(self.anim.names, self.anim.parent, self.anim.offsets, joint_reduction)
# Downsample to 30 fps for our application
if self.frametime < 0.0084:
self.frametime *= 2
self.anim.positions = self.anim.positions[::2]
self.anim.rotations = self.anim.rotations[::2]
if self.frametime < 0.017:
self.frametime *= 2
self.anim.positions = self.anim.positions[::2]
self.anim.rotations = self.anim.rotations[::2]
# Scale by 1/100 if it's raw exported bvh from blender
if not no_scale and self.skeleton.offsets[0, 1] > 10:
self.scale(1. / 100)
# Scale by 0.14 for Xia's dataset
if not no_scale and self.skeleton.skeleton_type == 3:
self.scale(0.14)
self.requires_contact = requires_contact
if requires_contact:
self.contact_names = self.skeleton.contact_names
else:
self.contact_names = []
self.fk = ForwardKinematicsJoint(self.skeleton.parent, self.skeleton.offsets)
self.writer = WriterWrapper(self.skeleton.parent, self.skeleton.offsets)
if self.requires_contact:
gl_pos = self.joint_position()
self.contact_label = foot_contact(gl_pos[:, self.skeleton.contact_id],
threshold=self.skeleton.contact_threshold)
self.gl_pos = gl_pos
def local_pos(self):
gl_pos = self.joint_position()
local_pos = gl_pos - gl_pos[:, 0:1, :]
return local_pos[:, 1:]
def scale(self, ratio):
self.anim.offsets *= ratio
self.anim.positions *= ratio
def to_tensor(self, repr='euler', rot_only=False):
if repr not in ['euler', 'quat', 'quaternion', 'repr6d']:
raise Exception('Unknown rotation representation')
positions = self.get_position()
rotations = self.get_rotation(repr=repr)
if rot_only:
return rotations.reshape(rotations.shape[0], -1)
if self.requires_contact:
virtual_contact = torch.zeros_like(rotations[:, :len(self.skeleton.contact_id)])
virtual_contact[..., 0] = self.contact_label
rotations = torch.cat([rotations, virtual_contact], dim=1)
rotations = rotations.reshape(rotations.shape[0], -1)
return torch.cat((rotations, positions), dim=-1)
def joint_position(self):
positions = torch.tensor(self.anim.positions[:, 0, :], dtype=torch.float)
rotations = self.anim.rotations[:, self.skeleton.corps, :]
rotations = Quaternions.from_euler(np.radians(rotations)).qs
rotations = torch.tensor(rotations, dtype=torch.float)
j_loc = self.fk.forward(rotations, positions)
return j_loc
def get_rotation(self, repr='quat'):
rotations = self.anim.rotations[:, self.skeleton.corps, :]
if repr == 'quaternion' or repr == 'quat' or repr == 'repr6d':
rotations = Quaternions.from_euler(np.radians(rotations)).qs
rotations = torch.tensor(rotations, dtype=torch.float)
if repr == 'repr6d':
rotations = quat2repr6d(rotations)
if repr == 'euler':
rotations = torch.tensor(rotations, dtype=torch.float)
return rotations
def get_position(self):
return torch.tensor(self.anim.positions[:, 0, :], dtype=torch.float)
def dfs(self, x, vis, dist):
fa = self.skeleton.parent
vis[x] = 1
for y in range(len(fa)):
if (fa[y] == x or fa[x] == y) and vis[y] == 0:
dist[y] = dist[x] + 1
self.dfs(y, vis, dist)
def get_neighbor(self, threshold, enforce_contact=False):
fa = self.skeleton.parent
neighbor_list = []
for x in range(0, len(fa)):
vis = [0 for _ in range(len(fa))]
dist = [0 for _ in range(len(fa))]
self.dfs(x, vis, dist)
neighbor = []
for j in range(0, len(fa)):
if dist[j] <= threshold:
neighbor.append(j)
neighbor_list.append(neighbor)
contact_list = []
if self.requires_contact:
for i, p_id in enumerate(self.skeleton.contact_id):
v_id = len(neighbor_list)
neighbor_list[p_id].append(v_id)
neighbor_list.append(neighbor_list[p_id])
contact_list.append(v_id)
root_neighbor = neighbor_list[0]
id_root = len(neighbor_list)
if enforce_contact:
root_neighbor = root_neighbor + contact_list
for j in contact_list:
neighbor_list[j] = list(set(neighbor_list[j]))
root_neighbor = list(set(root_neighbor))
for j in root_neighbor:
neighbor_list[j].append(id_root)
root_neighbor.append(id_root)
neighbor_list.append(root_neighbor) # Neighbor for root position
return neighbor_list
|