smaller error but still not same

smpl_tf.py

@@ -84,12 +84,13 @@ class SMPL:
         self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
 
         self.weights_all = self.weights
-        if use_one_hot:
+        # if use_one_hot:
-            self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
+        #     self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
-            self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1])
+        #     self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1])
-        else:
+        # else:
-            self.weights = tf.gather(self.weights, self.to_keep, axis=1)
+        #     self.weights = tf.gather(self.weights, self.to_keep, axis=1)
         
+
         self.shape = (1,27578,3)
 
         if orig_smpl:
@@ -189,21 +190,21 @@ class SMPL:
             v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
             v = tf.reshape(v, (-1, 4))[:, :3]
 
-            return v, tf.squeeze(J), v_posed
+            return v, tf.squeeze(J), v_posed, lrotmin
 
         gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)
-        T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
+        T_out, J_out, v_posed,lrotmin = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, 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)
-        v_posed = tf.subtract(v_posed, J0)
+        # v_posed = tf.subtract(v_posed, J0)
         
         if remove:
             T_out = tf.gather(T_out, self.to_keep, axis=1)
-            v_posed = tf.gather(v_posed, self.to_keep, axis=1)
+            # v_posed = tf.gather(v_posed, self.to_keep, axis=1)
         
-        return T_out, J_out, v_posed
+        return T_out, J_out, lrotmin, v_posed
 
 
     def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):

smpl_torch.py

@@ -10,10 +10,10 @@ from time import time
 from torch.nn import Module
 from utils import plot_mesh
 
-class SMPLModel(Module):
+class SMPL(Module):
   def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
     
-    super(SMPLModel, self).__init__()
+    super(SMPL, self).__init__()
     self.J_regressor = []
     self.v_template = []
     self.shapedirs = []
@@ -80,11 +80,14 @@ class SMPLModel(Module):
     -------
     Rotation matrix of shape [batch_size * angle_num, 3, 3].
     """
-    eps = r.clone().normal_(std=1e-8)
+    # eps = r.clone().normal_(std=1e-8)
-    theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    theta = torch.norm(r, dim=(1, 2), keepdim=True)
+    theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps)
+
     theta_dim = theta.shape[0]
     r_hat = r / theta
-    cos = torch.cos(theta)
+    # cos = torch.cos(theta)
     z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device)
     m = torch.stack(
       (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
@@ -93,8 +96,9 @@ class SMPLModel(Module):
     i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \
              + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device)
     A = r_hat.permute(0, 2, 1)
-    dot = torch.matmul(A, r_hat)
+    # dot = torch.matmul(A, r_hat)
-    R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta))
     return R
 
   @staticmethod
@@ -138,7 +142,7 @@ class SMPLModel(Module):
       for f in self.faces + 1:
         fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
 
-  def forward(self, betas, pose, trans, remove=True):
+  def forward(self, betas, pose, remove=True):
     
     """
           Construct a compute graph that takes in parameters and outputs a tensor as
@@ -211,7 +215,7 @@ class SMPLModel(Module):
     )
     v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
     v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
-    result = v + torch.reshape(trans, (batch_num, 1, 3))
+    result = v
     # estimate 3D joint locations
     # print(result.shape)
     # print(self.joint_regressor.shape)
@@ -224,7 +228,7 @@ class SMPLModel(Module):
     if remove:
       result = result[:,self.to_keep,:]
 
-    return result, joints
+    return result, joints, lrotmin, v_posed
 
 
 def test_gpu(gpu_id=[0], path=''):
@@ -239,7 +243,7 @@ def test_gpu(gpu_id=[0], path=''):
   beta_size = 10
 
   np.random.seed(9608)
-  model = SMPLModel(device=device, path=path)
+  model = SMPL(device=device, path=path)
   for i in range(5):
       pose = torch.from_numpy((np.random.rand(4, pose_size) - 0.5) * 0.4)\
               .type(torch.float32).to(device)
@@ -247,8 +251,7 @@ def test_gpu(gpu_id=[0], path=''):
               .type(torch.float32).to(device)
       betas = torch.concat((betas, torch.zeros((4,1), device=device)), axis=1)
       s = time()
-      trans = torch.from_numpy(np.zeros((4, 3))).type(torch.float32).to(device)
+      result, joints = model(betas, pose)
-      result, joints = model(betas, pose, trans)
       print(time() - s)
       plot_mesh(result[0].cpu().numpy(), model.f).show()
       

test.py

@@ -0,0 +1,52 @@
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '5'
+
+import numpy as np
+from smpl_tf import SMPL as SMPL_tf
+from smpl_torch import SMPL as SMPL_torch
+from utils import plot_mesh, plot_mesh_2
+from time import time
+
+import torch
+import tensorflow as tf
+
+if __name__ == '__main__':
+    device = torch.device('cuda')
+
+    smpl_tf : SMPL_tf = SMPL_tf('model_files')
+    smpl_torch : SMPL_torch = SMPL_torch(device=device, path='model_files')
+
+
+    pose = ((np.random.rand(4, 72) - 0.5) * 0.4).astype(np.float32)
+    betas = ((np.random.rand(4, 10) - 0.5) * 0.06).astype(np.float32)
+    betas = np.concatenate((betas, np.array([0,1,0,1], dtype=np.float32)[:,None]), axis=1)
+
+    # pose = np.zeros((4, 72), dtype=np.float32)
+    # betas = np.zeros((4, 11), dtype=np.float32)
+    
+    pose_torch = torch.tensor(pose, device=device)
+    betas_torch = torch.tensor(betas, device=device)
+
+    result_torch, joints_torch, _, _ = smpl_torch(betas_torch, pose_torch)
+    s = time()
+    result_torch, joints_torch, v_shaped_1, v_posed_1 = smpl_torch(betas_torch, pose_torch)
+    print(time() - s)
+    R_1 = smpl_torch.rodrigues(pose_torch.view(-1, 1, 3)).reshape(4, -1, 3, 3)
+
+    # plot_mesh(result[0].cpu().numpy(), model.f).show()
+
+    R = smpl_tf.run_rodrigues(pose)
+    result_tf, joints_torch, _, _ = smpl_tf.run(R, betas)
+    s = time()
+    R = smpl_tf.run_rodrigues(pose)
+    result_tf, joints_torch, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
+    print(time() - s)
+
+    result_torch = result_torch.cpu().numpy()
+    result_tf = result_tf.numpy()
+    diff = np.abs(result_torch - result_tf).sum(-1).sum(-1)
+
+    plot_mesh_2(result_torch[0], result_tf[0], smpl_tf.f).show()
+    plot_mesh_2(result_torch[1], result_tf[1], smpl_tf.f).show()
+    pass

utils.py

@@ -1,4 +1,5 @@
 import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
 
 def plot_mesh(verts, faces):
@@ -25,4 +26,43 @@ def plot_mesh(verts, faces):
     
     fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0))
 
+    return fig
+
+
+def plot_mesh_2(verts_1, verts_2, faces):
+    # fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
+    fig = go.Figure()
+    
+    fig.add_trace(go.Mesh3d(
+            x=verts_1[:,0],
+            y=verts_1[:,2],
+            z=verts_1[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="",            
+        ),)
+
+    fig.add_trace(go.Mesh3d(
+            x=verts_2[:,0],
+            y=verts_2[:,2],
+            z=verts_2[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="",            
+        ),)
+
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    # scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1)
+                    )
+    helper = dict(xaxis = dict(range=[-1,1], ),
+                yaxis = dict(range=[-1,1], ),
+                zaxis = dict(range=[-1,1], ) )
+    fig.update_layout(scene = helper, 
+                        # scene2 = helper
+    )
+    
+    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0))
+
     return fig