Skip to content

Commit

Permalink
remove unnecessary transpose
Browse files Browse the repository at this point in the history
  • Loading branch information
CalciferZh committed Mar 12, 2018
1 parent 506b7ab commit 6cab02e
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 172 deletions.
94 changes: 32 additions & 62 deletions smpl_np.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ def rodrigues_single(r):
theta = np.linalg.norm(r)
if theta == 0:
return np.eye(3)
r_hat = (r / theta).reshape(3, 1)
r_hat = np.expand_dims((r / theta), axis=1)
cos = np.cos(theta)
m = np.array([
0, -r_hat[2], r_hat[1], r_hat[2], 0, -r_hat[0], -r_hat[1], r_hat[0], 0
Expand All @@ -17,18 +17,17 @@ def rodrigues_single(r):


def rodrigues(r):
theta = np.linalg.norm(r, axis=(0, 1))
r_hat = r / theta
theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
r_hat = r / (theta + 1e-8)
cos = np.cos(theta) # theta is radius
z_stick = np.zeros(theta.shape[0])
m = np.hstack((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)).reshape(
(3, 3, -1))
i_cube = np.broadcast_to(
np.eye(3).reshape((3, 3, 1)), (3, 3, theta.shape[0]))
A = np.transpose(r_hat, axes=(2, 1, 0))
B = np.transpose(r_hat, axes=(2, 0, 1))
dot = np.matmul(A, B).transpose((1, 2, 0))
m = np.dstack((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)).reshape(
(-1, 3, 3))
i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), (theta.shape[0], 3, 3))
A = np.transpose(r_hat, axes=(0, 2, 1))
B = r_hat
dot = np.matmul(A, B)
R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m
return R

Expand All @@ -38,8 +37,7 @@ def with_zeros(x):


def pack(x):
depth = x.shape[-1]
return np.hstack((np.zeros((4, 3, depth)), x.reshape(4, 1, depth)))
return np.dstack((np.zeros((x.shape[0], 4, 3)), x))


def smpl_model(model_path, betas, pose, trans):
Expand All @@ -59,78 +57,50 @@ def smpl_model(model_path, betas, pose, trans):
i: id_to_col[kintree_table[0, i]]
for i in range(1, kintree_table.shape[1])
}

v_shaped = shapedirs.dot(betas) + v_template

J = J_regressor.dot(v_shaped)

pose_cube = pose[3:].reshape((-1, 1, 3)).transpose((1, 2, 0))

pose_cube = pose[3:].reshape((-1, 1, 3))
R_cube = rodrigues(pose_cube)
I_cube = np.broadcast_to(
np.eye(3).reshape((3, 3, 1)), (3, 3, R_cube.shape[2]))
lrotmin = (R_cube - I_cube).transpose(2, 0, 1).ravel()

I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), (R_cube.shape[0], 3, 3))
lrotmin = (R_cube - I_cube).ravel()
v_posed = v_shaped + posedirs.dot(lrotmin)

pose = pose.reshape((-1, 3))

depth = kintree_table.shape[1]

results = np.zeros((4, 4, depth))

results[:, :, 0] = with_zeros(
np.hstack((rodrigues_single(pose[0, :]), J[0, :].reshape((3, 1)))))

results = np.empty((kintree_table.shape[1], 4, 4))
results[0, :, :] = with_zeros(np.hstack((rodrigues_single(pose[0, :]), J[0, :].reshape((3, 1)))))
for i in range(1, kintree_table.shape[1]):
results[:, :, i] = results[:, :, parent[i]].dot(
with_zeros(
np.hstack((rodrigues_single(pose[i, :]),
((J[i, :] - J[parent[i], :]).reshape((3, 1)))))))


results = results - pack(
np.matmul(
results.transpose((2, 0, 1)),
np.hstack(
(
J,
np.zeros((24, 1))
)
).reshape((24, 4, 1)).transpose((0, 1, 2))
).transpose((1, 2, 0))
)

T = results.dot(weights.T)

rest_shape_h = np.vstack((v_posed.T, np.ones((1, v_posed.shape[0]))))


v = np.matmul(
T.transpose((2, 0, 1)),
rest_shape_h.reshape((4, 1, -1)).transpose((2, 0, 1))).reshape(
(-1, 4))[:, :3]

results[i, :, :] = results[parent[i], :, :].dot(with_zeros(np.hstack((rodrigues_single(pose[i, :]), ((J[i, :] - J[parent[i], :]).reshape((3, 1)))))))
results = results - pack(np.matmul(results, np.hstack((J, np.zeros((24, 1)))).reshape((24, 4, 1))))
T = np.tensordot(weights, results, axes=((1), (0)))
rest_shape_h = np.hstack((v_posed, np.ones((v_posed.shape[0], 1))))
v = np.matmul(T, rest_shape_h.reshape((-1, 4, 1))).reshape((-1, 4))[:, :3]
result = v + trans.reshape((1, 3))

return result, f


if __name__ == '__main__':
pose_size = 72
beta_size = 10
trans = np.zeros(3)

np.random.seed(9608)
pose = (np.random.rand(pose_size) - 0.5) * 0.4
# pose[0:3] = np.array([0, np.pi, 0])
betas = (np.random.rand(beta_size) - 0.5) * 0.06
# pose = np.zeros(pose_size)
# betas = np.zeros(beta_size)
trans = np.zeros(3)

# scale = 450
# trans = np.array([0, 0.3, 5.5])

result, faces = smpl_model('./model.pkl', betas, pose, trans)

outmesh_path = './hello_smpl.obj'
# result *= scale

outmesh_path = './smpl.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: # Faces are 1-based, not 0-based in obj files
for f in faces + 1:
fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
143 changes: 33 additions & 110 deletions smpl_tf.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,55 +8,51 @@ def rodrigues_single(r):
theta = tf.maximum(theta, 1e-8)
r_hat = tf.reshape(r / theta, (3, 1))
cos = tf.cos(theta)
z = tf.zeros((1), dtype=tf.float64)
z = tf.zeros((1), dtype=tf.float32)
m = tf.reshape(
tf.stack((z, -r_hat[2], r_hat[1], r_hat[2], z, -r_hat[0], -r_hat[1],
r_hat[0], z)), (3, 3))
R = cos * tf.eye(3, dtype=tf.float64) + (
R = cos * tf.eye(3, dtype=tf.float32) + (
1 - cos) * r_hat * tf.transpose(r_hat) + tf.sin(theta) * m
return R


def rodrigues(r):
theta = tf.norm(r, axis=(0, 1))
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()[0], dtype=tf.float64) # simply use static shape here
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))
m = tf.reshape(m, (3, 3, -1))
i_cube = tf.reshape(tf.eye(3, dtype=tf.float64), (3, 3, 1)) + tf.zeros(
(3, 3, theta.get_shape()[0]), dtype=tf.float64)
A = tf.transpose(r_hat, (2, 1, 0))
B = tf.transpose(r_hat, (2, 0, 1))
dot = tf.transpose(tf.matmul(A, B), (1, 2, 0))
(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.float64)), axis=0)
return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0)


def pack(x):
depth = x.get_shape().as_list()[-1]
z_cube = tf.zeros((4, 3, depth), dtype=tf.float64)
x_stick = tf.reshape(x, (4, 1, depth))
return tf.concat((z_cube, x_stick), axis=1)
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.float64))
weights = tf.constant(params['weights'], dtype=np.float64)
posedirs = tf.constant(params['posedirs'], dtype=np.float64)
v_template = tf.constant(params['v_template'], dtype=np.float64)
shapedirs = tf.constant(params['shapedirs'], dtype=np.float64)
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']
Expand All @@ -65,97 +61,24 @@ def smpl_model(model_path, betas, pose, trans):
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.transpose(tf.reshape(pose[3:], (-1, 1, 3)), (1, 2, 0))

pose_cube = tf.reshape(pose[3:], (-1, 1, 3))
R_cube = rodrigues(pose_cube)
I_cube = tf.reshape(tf.eye(3, dtype=tf.float64), (3, 3, 1)) + tf.zeros(
(3, 3, R_cube.get_shape()[2]), dtype=tf.float64)
lrotmin = tf.squeeze(
tf.reshape(tf.transpose(R_cube - I_cube, (2, 0, 1)), [1, -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))

depth = kintree_table.shape[1]

results = []

results.append(
with_zeros(
tf.concat(
(
rodrigues_single(pose[0, :]),
tf.reshape(
J[0, :],
(3, 1)
)
),
axis=1
)
)
)

results.append(with_zeros(tf.concat((rodrigues_single(pose[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(
(
rodrigues_single(pose[i, :]),
tf.reshape(J[i, :] - J[parent[i], :], (3, 1))
),
axis=1
)
)
)
)

stacked = tf.stack(results, axis=2)

results = stacked - pack(
tf.transpose(
tf.matmul(
tf.transpose(stacked, (2, 0, 1)),
tf.reshape(
tf.concat(
(
J,
tf.zeros((24, 1), dtype=tf.float64)
),
axis=1
),
(24, 4, 1)
)
),
(1, 2, 0)
)
)

T = tf.tensordot(results, weights, axes=[[2], [1]])

rest_shape_h = tf.concat((tf.transpose(v_posed),
tf.ones((1, v_posed.get_shape()[0]), dtype=tf.float64)), axis=0)

v = tf.matmul(
tf.transpose(T, (2, 0, 1)),
tf.transpose(
tf.reshape(
rest_shape_h,
(4, 1, -1)
),
(2, 0, 1)
)
)

results.append(tf.matmul(results[parent[i]], with_zeros(tf.concat((rodrigues_single(pose[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

Expand All @@ -169,17 +92,17 @@ def smpl_model(model_path, betas, pose, trans):
betas = (np.random.rand(beta_size) - 0.5) * 0.06
trans = np.zeros(3)

pose = tf.constant(pose, dtype=tf.float64)
betas = tf.constant(betas, dtype=tf.float64)
trans = tf.constant(trans, dtype=tf.float64)
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 = './hello_smpl.obj'
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]))
Expand Down

0 comments on commit 6cab02e

Please sign in to comment.