smaller error but still not same

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)
-    theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    # 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, 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)
-    R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    # dot = torch.matmul(A, r_hat)
+    # 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, trans)
+      result, joints = model(betas, pose)
       print(time() - s)
       plot_mesh(result[0].cpu().numpy(), model.f).show()