From 74a0d58b165d2abeaac75b37791cdd2bf4369ad1 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:11:33 +0000 Subject: [PATCH] orig smpl code for reference --- smpl_orig.py | 178 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 178 insertions(+) create mode 100644 smpl_orig.py diff --git a/smpl_orig.py b/smpl_orig.py new file mode 100644 index 0000000..02daafc --- /dev/null +++ b/smpl_orig.py @@ -0,0 +1,178 @@ +import tensorflow as tf +import numpy as np +import pickle + + +def rodrigues(r): + """ + Rodrigues' rotation formula that turns axis-angle tensor into rotation + matrix in a batch-ed manner. + Parameter: + ---------- + r: Axis-angle rotation tensor of shape [batch_size, 1, 3]. + Return: + ------- + Rotation matrix of shape [batch_size, 3, 3]. + """ + theta = tf.norm(r + tf.random_normal(r.shape, 0, 1e-8, dtype=tf.float64), axis=(1, 2), keepdims=True) + # avoid divide by zero + r_hat = r / theta + cos = tf.cos(theta) + z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float64) + 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.float64), axis=0) + tf.zeros( + (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float64) + 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): + """ + Append a [0, 0, 0, 1] tensor to a [3, 4] tensor. + Parameter: + --------- + x: Tensor to be appended. + Return: + ------ + Tensor after appending of shape [4,4] + """ + ret = tf.concat( + (x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float64)), + axis=0 + ) + return ret + + +def pack(x): + """ + Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor. + Parameter: + ---------- + x: A tensor of shape [batch_size, 4, 1] + Return: + ------ + A tensor of shape [batch_size, 4, 4] after appending. + """ + ret = tf.concat( + (tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float64), x), + axis=2 + ) + return ret + + +def smpl_model(model_path, betas, pose, trans, simplify=False): + """ + Construct a compute graph that takes in parameters and outputs a tensor as + model vertices. Face indices are also returned as a numpy ndarray. + Parameters: + --------- + pose: Also known as 'theta', a [24,3] tensor indicating child joint rotation + relative to parent joint. For root joint it's global orientation. + Represented in a axis-angle format. + betas: Parameter for model shape. A tensor of shape [10] as coefficients of + PCA components. Only 10 components were released by SMPL author. + trans: Global translation tensor of shape [3]. + Return: + ------ + A tensor for vertices, and a numpy ndarray as face indices. + """ + # For detailed comments see smpl_np.py + 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) + 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) + if simplify: + v_posed = v_shaped + else: + R_cube = R_cube_big[1:] + I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float64), axis=0) + \ + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float64) + lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1))) + v_posed = v_shaped + tf.tensordot(posedirs, lrotmin, axes=[[2], [0]]) + 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.float64)), 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.float64)), + 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.zeros(pose_size) + betas = np.zeros(beta_size) + 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) + + output, faces = smpl_model('./model.pkl', betas, pose, trans, True) + 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])) \ No newline at end of file