add comments

README.md

@@ -5,19 +5,12 @@ Numpy and Tensorflow implementation of SMPL model. For any questions, feel free
 
 I wrote this because the author-provided implementation was mainly based on [chumpy](https://github.com/mattloper/chumpy) in Python 2, which is kind of unpopular. Meanwhile, the official one cannot run on GPU.
 
-This numpy version is faster(since some computation is re-wrote in a vectorized manner) and easier to understand(hope so), and the tensorflow version can run on GPU.
+This numpy version is faster (since some computation is re-wrote in a vectorized manner) and easier to understand (hope so), and the tensorflow version can run on GPU.
 
 For more details about SMPL model, see [SMPL](http://smpl.is.tue.mpg.de/).
 
-Also, I provide a file `CMU_Mocap_Markers.pp`, which gives the correspondence between SMPL model and [CMU Mocap Dataset](http://mocap.cs.cmu.edu/) markers in .c3d files. For more details see the Usage section.
-
 ## Usage
 
 1. Download the model file [here](http://smpl.is.tue.mpg.de/downloads).
-2. Run `python preprocess.py /PATH/TO/THE/DOWNLOADED/MODEL` to preprocess the official model. `preprocess.py` will create a new file `model.pkl`. `smpl_np.py` and `smpl_tf.py` both rely on `model.pkl`. **NOTE: the official pickle model contains `chumpy` object, so `prerocess.py` requires `chumpy` to extract official model.** Actually you need to modify chumpy's cource code a little bit to make it compatible to `preprocess.py`.
-3. Run `python smpl_np.py` to see the example.
-4. About `CMU_Mocap_Markers.pp`: you can first generate a standard SMPL model mesh(zero pose and zero beta), open it in MeshLab, and load this file in MeshLab. It gives 42 markers' position on the model surface. I simply mark these things by hand so there might be some small errors.
-
-## One More Thing
-
-If this repo is used in any publication or project, it would be nice to let me know. I will be very happy and encouraged =)
+2. Run `python preprocess.py /PATH/TO/THE/DOWNLOADED/MODEL` to preprocess the official model. `preprocess.py` will create a new file `model.pkl`. `smpl_np.py` and `smpl_tf.py` both rely on `model.pkl`. **NOTE**: the official pickle model contains `chumpy` object, so `prerocess.py` requires `chumpy` to extract official model. You need to modify chumpy's cource code a little bit to make it compatible to `preprocess.py` (and Python 3).
+3. Run `python smpl_np.py` or `python smpl_tf.py` to see the example.

smpl_np.py

@@ -3,122 +3,211 @@
 
 
 class SMPLModel():
-    def __init__(self, model_path):
-        with open(model_path, 'rb') as f:
-            params = pickle.load(f)
-
-            self.J_regressor = params['J_regressor']
-            self.weights = params['weights']
-            self.posedirs = params['posedirs']
-            self.v_template = params['v_template']
-            self.shapedirs = params['shapedirs']
-            self.faces = params['f']
-            self.kintree_table = params['kintree_table']
-
-        id_to_col = {self.kintree_table[1, i]: i for i in range(self.kintree_table.shape[1])}
-        self.parent = {
-            i: id_to_col[self.kintree_table[0, i]]
-            for i in range(1, self.kintree_table.shape[1])
-        }
-
-        self.pose_shape = [24, 3]
-        self.beta_shape = [10]
-        self.trans_shape = [3]
-
-        self.pose = np.zeros(self.pose_shape)
-        self.beta = np.zeros(self.beta_shape)
-        self.trans = np.zeros(self.trans_shape)
-
-        self.verts = None
-        self.J = None
-        self.R = None
-
-        self.update()
-
-    def set_params(self, pose=None, beta=None, trans=None):
-        if pose is not None:
-            self.pose = pose
-        if beta is not None:
-            self.beta = beta
-        if trans is not None:
-            self.trans = trans
-        self.update()
-        return self.verts
-
-    def update(self):
-        v_shaped = self.shapedirs.dot(self.beta) + self.v_template # how beta affect body shape
-        self.J = self.J_regressor.dot(v_shaped) # joints location
-        pose_cube = self.pose.reshape((-1, 1, 3))
-        self.R = self.rodrigues(pose_cube) # rotation matrix for each joint
-        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), (self.R.shape[0]-1, 3, 3))
-        lrotmin = (self.R[1:] - I_cube).ravel()
-        v_posed = v_shaped + self.posedirs.dot(lrotmin) # how pose affect body shape in zero pose
-        G = np.empty((self.kintree_table.shape[1], 4, 4)) # world transformation of each joint
-        G[0] = self.with_zeros(np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
-        for i in range(1, self.kintree_table.shape[1]):
-            G[i] = G[self.parent[i]].dot(
-                self.with_zeros(
-                    np.hstack(
-                        [self.R[i], ((self.J[i, :] - self.J[self.parent[i], :]).reshape([3, 1]))]
-                    )
-                )
-            )
-        # remove the transformation due to the rest pose
-        G = G - self.pack(
-            np.matmul(
-                G,
-                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])
-                )
-            )
-        T = np.tensordot(self.weights, G, axes=[[1], [0]]) # transformation of each vertex
-        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]
-        self.verts = v + self.trans.reshape([1, 3])
-
-    def rodrigues(self, r):
-        theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
-        # avoid zero divide
-        theta = np.maximum(theta, np.finfo(np.float64).tiny)
-        r_hat = r / theta
-        cos = np.cos(theta)
-        z_stick = np.zeros(theta.shape[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]
+  def __init__(self, model_path):
+    """
+    SMPL model.
+
+    Parameter:
+    ---------
+    model_path: Path to the SMPL model parameters, pre-processed by
+    `preprocess.py`.
+
+    """
+    with open(model_path, 'rb') as f:
+      params = pickle.load(f)
+
+      self.J_regressor = params['J_regressor']
+      self.weights = params['weights']
+      self.posedirs = params['posedirs']
+      self.v_template = params['v_template']
+      self.shapedirs = params['shapedirs']
+      self.faces = params['f']
+      self.kintree_table = params['kintree_table']
+
+    id_to_col = {
+      self.kintree_table[1, i]: i for i in range(self.kintree_table.shape[1])
+    }
+    self.parent = {
+      i: id_to_col[self.kintree_table[0, i]]
+      for i in range(1, self.kintree_table.shape[1])
+    }
+
+    self.pose_shape = [24, 3]
+    self.beta_shape = [10]
+    self.trans_shape = [3]
+
+    self.pose = np.zeros(self.pose_shape)
+    self.beta = np.zeros(self.beta_shape)
+    self.trans = np.zeros(self.trans_shape)
+
+    self.verts = None
+    self.J = None
+    self.R = None
+
+    self.update()
+
+  def set_params(self, pose=None, beta=None, trans=None):
+    """
+    Set pose, shape, and/or translation parameters of SMPL model. Verices of the
+    model will be updated and returned.
+
+    Prameters:
+    ---------
+    pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation
+    relative to parent joint. For root joint it's global orientation.
+    Represented in a axis-angle format.
+
+    beta: Parameter for model shape. A vector of shape [10]. Coefficients for
+    PCA component. Only 10 components were released by MPI.
+
+    trans: Global translation of shape [3].
+
+    Return:
+    ------
+    Updated vertices.
+
+    """
+    if pose is not None:
+      self.pose = pose
+    if beta is not None:
+      self.beta = beta
+    if trans is not None:
+      self.trans = trans
+    self.update()
+    return self.verts
+
+  def update(self):
+    """
+    Called automatically when parameters are updated.
+
+    """
+    # how beta affect body shape
+    v_shaped = self.shapedirs.dot(self.beta) + self.v_template
+    # joints location
+    self.J = self.J_regressor.dot(v_shaped)
+    pose_cube = self.pose.reshape((-1, 1, 3))
+    # rotation matrix for each joint
+    self.R = self.rodrigues(pose_cube)
+    I_cube = np.broadcast_to(
+      np.expand_dims(np.eye(3), axis=0),
+      (self.R.shape[0]-1, 3, 3)
+    )
+    lrotmin = (self.R[1:] - I_cube).ravel()
+    # how pose affect body shape in zero pose
+    v_posed = v_shaped + self.posedirs.dot(lrotmin)
+    # world transformation of each joint
+    G = np.empty((self.kintree_table.shape[1], 4, 4))
+    G[0] = self.with_zeros(np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+    for i in range(1, self.kintree_table.shape[1]):
+      G[i] = G[self.parent[i]].dot(
+        self.with_zeros(
+          np.hstack(
+            [self.R[i],((self.J[i, :]-self.J[self.parent[i],:]).reshape([3,1]))]
+          )
         )
-        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
-
-    def with_zeros(self, x):
-        return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))
-
-
-    def pack(self, x):
-        return np.dstack((np.zeros((x.shape[0], 4, 3)), x))
-
-    def save_to_obj(self, path):
-        with open(path, 'w') as fp:
-            for v in self.verts:
-                fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
-            for f in self.faces + 1:
-                fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
+      )
+    # remove the transformation due to the rest pose
+    G = G - self.pack(
+      np.matmul(
+        G,
+        np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])
+        )
+      )
+    # transformation of each vertex
+    T = np.tensordot(self.weights, G, 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]
+    self.verts = v + self.trans.reshape([1, 3])
+
+  def rodrigues(self, r):
+    """
+    Rodrigues' rotation formula that turns axis-angle vector into rotation
+    matrix in a batch-ed manner.
+
+    Parameter:
+    ----------
+    r: Axis-angle rotation vector of shape [batch_size, 1, 3].
+
+    Return:
+    -------
+    Rotation matrix of shape [batch_size, 3, 3].
+
+    """
+    theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
+    # avoid zero divide
+    theta = np.maximum(theta, np.finfo(np.float64).tiny)
+    r_hat = r / theta
+    cos = np.cos(theta)
+    z_stick = np.zeros(theta.shape[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
+
+  def with_zeros(self, x):
+    """
+    Append a [0, 0, 0, 1] vector to a [3, 4] matrix.
+
+    Parameter:
+    ---------
+    x: Matrix to be appended.
+
+    Return:
+    ------
+    Matrix after appending of shape [4,4]
+
+    """
+    return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))
+
+  def pack(self, x):
+    """
+    Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched
+    manner.
+
+    Parameter:
+    ----------
+    x: Matrices to be appended of shape [batch_size, 4, 1]
+
+    Return:
+    ------
+    Matrix of shape [batch_size, 4, 4] after appending.
+
+    """
+    return np.dstack((np.zeros((x.shape[0], 4, 3)), x))
+
+  def save_to_obj(self, path):
+    """
+    Save the SMPL model into .obj file.
+
+    Parameter:
+    ---------
+    path: Path to save.
+
+    """
+    with open(path, 'w') as fp:
+      for v in self.verts:
+        fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+      for f in self.faces + 1:
+        fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
 
 
 
 if __name__ == '__main__':
-    smpl = SMPLModel('./model.pkl')
-    np.random.seed(9608)
-    pose = (np.random.rand(*smpl.pose_shape) - 0.5) * 0.4
-    beta = (np.random.rand(*smpl.beta_shape) - 0.5) * 0.06
-    trans = np.zeros(smpl.trans_shape)
-    smpl.set_params(beta=beta, pose=pose, trans=trans)
-    smpl.save_to_obj('./smpl_np.obj')
+  smpl = SMPLModel('./model.pkl')
+  np.random.seed(9608)
+  pose = (np.random.rand(*smpl.pose_shape) - 0.5) * 0.4
+  beta = (np.random.rand(*smpl.beta_shape) - 0.5) * 0.06
+  trans = np.zeros(smpl.trans_shape)
+  smpl.set_params(beta=beta, pose=pose, trans=trans)
+  smpl.save_to_obj('./smpl_np.obj')