add new options to smpl code

custom_utils/smpl.py

@@ -34,7 +34,7 @@ eval_joints = {
         
 class SMPL:
 
-    def __init__(self):
+    def __init__(self, orig_smpl=False, use_one_hot=False):
         self.J_regressor = []
         self.v_template = []
         self.shapedirs = []
@@ -82,7 +82,23 @@ class SMPL:
         self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
 
         self.weights_all = self.weights
-        self.weights = tf.gather(self.weights, self.to_keep, axis=1)
+        if use_one_hot:
+            self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
+            self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1])
+        else:
+            self.weights = tf.gather(self.weights, self.to_keep, axis=1)
+        
+        self.shape = (1,27578,3)
+
+        if orig_smpl:
+            self.to_keep = self.to_keep[self.to_keep < 6890]
+            self.J_regressor = self.J_regressor[:,:,:6890]
+            self.v_template = self.v_template[:,:6890,:]
+            self.shapedirs = self.shapedirs[:,:6890,:,:]
+            self.posedirs = self.posedirs[:,:6890,:,:]
+            self.weights_all = self.weights_all[:,:6890,:]
+            self.weights = tf.gather(self.weights_all, self.to_keep, axis=1)
+            self.shape = (1, 6890, 3)
 
     def Q2R(self, w, x, y, z):
         def diag(a, b):  # computes the diagonal entries,  1 - 2*a**2 - 2*b**2
@@ -141,8 +157,8 @@ class SMPL:
         #v_template = tf.gather(self.v_template, gender, axis=0)
         
         v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
+        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
+        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
         v_shaped = tf.stack(v_shaped, axis=-1)
 
         def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights_all, parent=self.parent):
@@ -187,11 +203,12 @@ class SMPL:
         
         return T_out, J_out, v_posed
 
+
     def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):
         
         v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
+        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
+        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
         v_shaped = tf.stack(v_shaped, axis=-1)
 
         if weights is None:
@@ -236,11 +253,14 @@ class SMPL:
 
         return T_out
 
-    def get_joints_template(self, betas):
+
-        v_shaped = []
+    def get_joints_template(self, betas, v_shaped=None):
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
+        
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
+        if v_shaped is None:
-        v_shaped = tf.stack(v_shaped, axis=-1)
+            v_shaped = []
+            v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
+            v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
+            v_shaped = tf.stack(v_shaped, axis=-1)
 
         def single_run(x, J_regressor=self.J_regressor):
             v_shaped = x[0]
@@ -254,9 +274,31 @@ class SMPL:
         
         return tf.subtract(J, J0)
 
+    def get_root_joint(self, betas, v_shaped=None):
+        
+        if v_shaped is None:
+            v_shaped = []
+            v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female
+            v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male
+            v_shaped = tf.stack(v_shaped, axis=-1)
 
-    def run_garment(self, pose, betas, v_cloth, filter=None):
+        def single_run(x, J_regressor=self.J_regressor):
-        J = self.get_joints_template(betas)            
+            v_shaped = x[0]
+            g = x[1]
+            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
+            return J
+
+        gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32)
+        J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32))
+        J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
+        
+        return J[:,0,:]
+
+        # return J0
+
+
+    def run_garment(self, pose, betas, v_cloth, v_shaped=None):
+        J = self.get_joints_template(betas, v_shaped)            
         R = self.run_rodrigues(pose)
         
         # if filter is not None:
@@ -269,7 +311,7 @@ class SMPL:
             v_cloth = x[3]
 
             R_cube_big = R
-            
+
             results = []
             results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
             for j in range(1, 24):
@@ -294,8 +336,59 @@ class SMPL:
         J_out = tf.subtract(J_out, J0)
         T_out = tf.subtract(T_out, J0)
         
+
         return T_out
 
+
+    def run_garment_orig(self, R, betas, v_cloth, J=None, weights=None):
+        batch_size = R.shape[0]
+        
+        if J is None:
+            J = self.get_joints_template(betas)
+        
+        if weights is None:
+            weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
+        # else:
+            # weights = tf.gather(weights, self.to_keep, axis=0)
+
+        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=weights, parent=self.parent):
+            R = x[0]
+            J = x[1]
+            g = x[2]
+            v_cloth = x[3]
+
+            R_cube_big = R
+            # 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)))
+            
+            results = []
+            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
+            for j in range(1, 24):
+                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
+            stacked = tf.stack(results, axis=0)
+            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
+            
+            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
+            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
+            
+            T = tf.tensordot(weights, results, axes=((1), (0)))
+            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.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]
+
+            return v, tf.squeeze(J)
+
+        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
+        T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32))
+        
+        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
+        J_out = tf.subtract(J_out, J0)
+        T_out = tf.subtract(T_out, J0)
+        
+        return T_out, J_out
+
+
     def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
         
         if J is None:
@@ -335,9 +428,14 @@ class SMPL:
         
         return T_out
 
+
     def run_backward(self, R, J, v_posed, weights):
         batch_size = R.get_shape()[0]
         
+        if weights.shape[0] != 14475:
+            weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
+
+
         J_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3
 
         J_out = []