orig smpl code for reference

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]))