-
Notifications
You must be signed in to change notification settings - Fork 140
/
Copy pathsmpl_tf.py
96 lines (79 loc) · 3.73 KB
/
smpl_tf.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
import tensorflow as tf
import numpy as np
import pickle
def rodrigues(r):
theta = tf.norm(r, axis=(1, 2), keepdims=True)
theta = tf.maximum(theta, 1e-8)
r_hat = r / theta
cos = tf.cos(theta)
z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
m = tf.stack(
(z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
-r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
m = tf.reshape(m, (-1, 3, 3))
i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
(theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
A = tf.transpose(r_hat, (0, 2, 1))
B = r_hat
dot = tf.matmul(A, B)
R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
return R
def with_zeros(x):
return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0)
def pack(x):
return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
def smpl_model(model_path, betas, pose, trans):
with open(model_path, 'rb') as f:
params = pickle.load(f)
J_regressor = tf.constant(np.array(params['J_regressor'].todense(), dtype=np.float32))
weights = tf.constant(params['weights'], dtype=np.float32)
posedirs = tf.constant(params['posedirs'], dtype=np.float32)
v_template = tf.constant(params['v_template'], dtype=np.float32)
shapedirs = tf.constant(params['shapedirs'], dtype=np.float32)
f = params['f']
kintree_table = params['kintree_table']
id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
parent = {
i: id_to_col[kintree_table[0, i]]
for i in range(1, kintree_table.shape[1])
}
v_shaped = tf.tensordot(shapedirs, betas, axes=[[2], [0]]) + v_template
J = tf.matmul(J_regressor, v_shaped)
pose_cube = tf.reshape(pose, (-1, 1, 3))
R_cube_big = rodrigues(pose_cube)
R_cube = R_cube_big[1:]
I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
v_posed = v_shaped + tf.tensordot(posedirs, lrotmin, axes=[[2], [0]])
pose = tf.reshape(pose, (-1, 3))
results = []
results.append(with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
for i in range(1, kintree_table.shape[1]):
results.append(tf.matmul(results[parent[i]], with_zeros(tf.concat((R_cube_big[i], tf.reshape(J[i, :] - J[parent[i], :], (3, 1))), axis=1))))
stacked = tf.stack(results, axis=0)
results = stacked - pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
T = tf.tensordot(weights, results, axes=((1), (0)))
rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
v = tf.reshape(v, (-1, 4))[:, :3]
result = v + tf.reshape(trans, (1, 3))
return result, f
if __name__ == '__main__':
pose_size = 72
beta_size = 10
np.random.seed(9608)
pose = (np.random.rand(pose_size) - 0.5) * 0.4
betas = (np.random.rand(beta_size) - 0.5) * 0.06
trans = np.zeros(3)
pose = tf.constant(pose, dtype=tf.float32)
betas = tf.constant(betas, dtype=tf.float32)
trans = tf.constant(trans, dtype=tf.float32)
output, faces = smpl_model('./model.pkl', betas, pose, trans)
sess = tf.Session()
result = sess.run(output)
outmesh_path = './smpl_tf.obj'
with open(outmesh_path, 'w') as fp:
for v in result:
fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
for f in faces + 1:
fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))