finalize

smpl_torch.py

@@ -134,13 +134,6 @@ class SMPL(Module):
     ret = torch.cat((zeros43, x), dim=3)
     return ret
 
-  def write_obj(self, verts, file_name):
-    with open(file_name, 'w') as fp:
-      for v in 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]))
 
   def forward(self, betas, pose, remove=True):
     
@@ -208,7 +201,9 @@ class SMPL(Module):
           )
         )
       )
-    # Restart from here
+    
+    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+            
     T = torch.stack([torch.tensordot(results[i], self.weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
     rest_shape_h = torch.cat(
       (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
@@ -216,11 +211,7 @@ class SMPL(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
-    # estimate 3D joint locations
-    # print(result.shape)
-    # print(self.joint_regressor.shape)
-    joints = torch.stack([torch.tensordot(result[i], self.J_regressor[g.to(torch.long)], dims=([0], [1])) for i, g in enumerate(betas[:,-1])]).transpose(1, 2)
-    
+
     root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
     result = result - root_joint
     joints = joints - root_joint
@@ -228,37 +219,4 @@ class SMPL(Module):
     if remove:
       result = result[:,self.to_keep,:]
 
-    return result, joints, lrotmin, v_posed
-
-
-def test_gpu(gpu_id=[0], path=''):
-  if len(gpu_id) > 0 and torch.cuda.is_available():
-    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id[0])
-    device = torch.device('cuda')
-  else:
-    device = torch.device('cpu')
-  #print(device)
-
-  pose_size = 72
-  beta_size = 10
-
-  np.random.seed(9608)
-  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)
-      betas = torch.from_numpy((np.random.rand(4, beta_size) - 0.5) * 0.06) \
-              .type(torch.float32).to(device)
-      betas = torch.concat((betas, torch.zeros((4,1), device=device)), axis=1)
-      s = time()
-      result, joints = model(betas, pose)
-      print(time() - s)
-      plot_mesh(result[0].cpu().numpy(), model.f).show()
-      
-  # outmesh_path = './smpl_torch_{}.obj'
-  # for i in range(result.shape[0]):
-      # model.write_obj(result[i], outmesh_path.format(i))
-  
-
-if __name__ == '__main__':
-  test_gpu([5], path='model_files')
+    return result, joints