found diff

smpl_tf.py

@@ -141,6 +141,7 @@ class SMPL:
     def pack(self, x):
         return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
         
+    # @tf.function
     def run_rodrigues(self, pose):
         batch_size = pose.shape[0]
         R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
@@ -151,6 +152,7 @@ class SMPL:
         #R_out = tf.stack(R_out, axis=0)
         return R_out
 
+    # @tf.function
     def run(self, R, betas, remove=True):
         batch_size = betas.shape[0]
         # R = self.run_rodrigues(pose)
@@ -182,14 +184,17 @@ class SMPL:
             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)
+            # 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[g], 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]
 
+            J = tf.transpose(tf.tensordot(v, self.J_regressor[g], axes=([0], [1])), [1,0])
+
             return v, tf.squeeze(J), v_posed, lrotmin
 
         gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)