1 from __future__
import absolute_import
2 from __future__
import division
3 from __future__
import print_function
5 import itertools, pkg_resources, sys
6 from distutils.version
import LooseVersion
7 if LooseVersion(pkg_resources.get_distribution(
"chainer").version) >= LooseVersion(
'7.0.0')
and \
8 sys.version_info.major == 2:
9 print(
'''Please install chainer < 7.0.0: 11 sudo pip install chainer==6.7.0 13 c.f https://github.com/jsk-ros-pkg/jsk_recognition/pull/2485 16 if [p
for p
in list(itertools.chain(*[pkg_resources.find_distributions(_)
for _
in sys.path]))
if "cupy-" in p.project_name ] == []:
17 print(
'''Please install CuPy 19 sudo pip install cupy-cuda[your cuda version] 21 sudo pip install cupy-cuda91 25 from chainer
import Parameter
26 from chainer
import Variable
28 import chainer.functions
as F
34 Computes absolute joint locations given pose. 36 rotate_base: if True, rotates the global rotation by 90 deg in x axis. 37 if False, this is the original SMPL coordinate. 40 Rs: N x 24 x 3 x 3 rotation vector of K joints 41 Js: N x 24 x 3, joint locations before posing 42 parent: 24 holding the parent id for each index 45 new_J : `Tensor`: N x 24 x 3 location of absolute joints 46 A : `Tensor`: N x 24 4 x 4 relative joint transformations for LBS. 51 print(
'Flipping the SMPL coordinate frame!!!!')
53 [[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=Rs.dtype)
54 rot_x = F.reshape(F.tile(rot_x, [N, 1]), [N, 3, 3])
55 root_rotation = F.matmul(Rs[:, 0, :, :], rot_x)
57 root_rotation = Rs[:, 0, :, :]
60 Js = F.expand_dims(Js, -1)
62 def make_A(R, t, name=None):
64 R_homo = F.pad(R, [[0, 0], [0, 1], [0, 0]],
'constant')
65 t_homo = F.concat([t, xp.ones([N, 1, 1],
'f')], 1)
66 return F.concat([R_homo, t_homo], 2)
68 A0 = make_A(root_rotation, Js[:, 0])
70 for i
in range(1, parent.shape[0]):
71 j_here = Js[:, i] - Js[:, parent[i]]
72 A_here = make_A(Rs[:, i], j_here)
74 results[parent[i]], A_here)
75 results.append(res_here)
78 results = F.stack(results, axis=1)
80 new_J = results[:, :, :3, 3]
86 Js_w0 = F.concat([Js, xp.zeros([N, 24, 1, 1],
'f')], 2)
87 init_bone = F.matmul(results, Js_w0)
89 init_bone = F.pad(init_bone, [[0, 0], [0, 0], [0, 0], [3, 0]],
'constant')
90 A = results - init_bone
97 vec is N x 3, batch_size is int 99 returns N x 3 x 3. Skew_sym version of each matrix. 102 if batch_size
is None:
103 batch_size = vec.shape[0]
104 col_inds = xp.array([1, 2, 3, 5, 6, 7])
106 F.repeat(col_inds.reshape(1, -1), batch_size, axis=0) +
107 F.repeat(F.reshape(xp.arange(0, batch_size) * 9, [-1, 1]), 6, axis=1),
111 [-vec[:, 2], vec[:, 1], vec[:, 2], -vec[:, 0], -vec[:, 1],
112 vec[:, 0]], axis=1), [-1])
113 res = Variable(xp.zeros((batch_size * 3 * 3),
'f'))
114 res.data[indices.reshape(-1).data] = updates.data
115 res = F.reshape(res, [batch_size, 3, 3])
123 batch_size = theta.shape[0]
126 angle = F.expand_dims(F.sqrt(F.batch_l2_norm_squared(theta + 1e-8)), -1)
127 r = F.expand_dims(theta / F.tile(angle, 3), -1)
129 angle = F.expand_dims(angle, -1)
132 cos = F.tile(cos, (3, 3))
133 sin = F.tile(sin, (3, 3))
135 outer = F.matmul(r, r, transb=
True)
137 eyes = F.tile(F.expand_dims(
138 Variable(xp.array(xp.eye(3),
'f')), 0), (batch_size, 1, 1))
139 R = cos * eyes + (1 - cos) * outer + sin *
batch_skew(r, batch_size)
148 4294967295, 0, 0, 0, 1,
149 2, 3, 4, 5, 6, 7, 8, 9, 9, 9,
150 12, 13, 14, 16, 17, 18, 19, 20, 21],
'i')
152 with self.init_scope():
154 self.
size = [self.v_template.shape[0], 3]
161 def __call__(self, beta, theta, get_skin=False, with_a=False):
162 batch_size = beta.shape[0]
167 v_shaped = F.reshape(
169 [-1, self.
size[0], self.
size[1]]) + \
170 F.repeat(self.
v_template[
None, ], batch_size, axis=0)
177 J = F.stack([Jx, Jy, Jz], axis=2)
187 pose_feature = F.reshape(Rs[:, 1:, :, :] -
188 F.repeat(F.repeat(Variable(self.xp.array(self.xp.eye(3),
'f'))[
189 None, ], 23, axis=0)[
None, ], batch_size, axis=0),
195 F.matmul(pose_feature, self.
posedirs),
196 [-1, self.
size[0], self.
size[1]]) + v_shaped
205 F.tile(self.
weights, (batch_size, 1)), [batch_size, -1, 24])
208 F.matmul(W, F.reshape(A, [batch_size, 24, 16])),
209 [batch_size, -1, 4, 4])
210 v_posed_homo = F.concat(
211 [v_posed, self.xp.ones([batch_size, v_posed.shape[1], 1],
'f')], 2)
212 v_homo = F.matmul(T, F.expand_dims(v_posed_homo, -1))
214 verts = v_homo[:, :, :3, 0]
220 joints = F.stack([joint_x, joint_y, joint_z], axis=2)
222 return verts, joints, Rs, A
def __call__(self, beta, theta, get_skin=False, with_a=False)
def batch_global_rigid_transformation(Rs, Js, parent, rotate_base=False)
def batch_rodrigues(theta)
def batch_skew(vec, batch_size=None)