optimize rot_6d instead of rot_mat

apps/infer.py

@@ -26,6 +26,7 @@ from lib.dataset.mesh_util import (
     unwrap,
     remesh,
     tensor2variable,
+    rot6d_to_rotmat
 )
 
 from lib.dataset.TestDataset import TestDataset
@@ -165,12 +166,16 @@ def generate_model(in_path, model_type):
         for _ in loop_smpl:
 
             optimizer_smpl.zero_grad()
+            
+            # 6d_rot to rot_mat
+            optimed_orient_mat = rot6d_to_rotmat(optimed_orient.view(-1,6)).unsqueeze(0)
+            optimed_pose_mat = rot6d_to_rotmat(optimed_pose.view(-1,6)).unsqueeze(0)
 
             if dataset_param["hps_type"] != "pixie":
                 smpl_out = dataset.smpl_model(
                     betas=optimed_betas,
-                    body_pose=optimed_pose,
-                    global_orient=optimed_orient,
+                    body_pose=optimed_pose_mat,
+                    global_orient=optimed_orient_mat,
                     pose2rot=False,
                 )
 
@@ -180,8 +185,8 @@ def generate_model(in_path, model_type):
                 smpl_verts, _, _ = dataset.smpl_model(
                     shape_params=optimed_betas,
                     expression_params=tensor2variable(data["exp"], device),
-                    body_pose=optimed_pose,
-                    global_pose=optimed_orient,
+                    body_pose=optimed_pose_mat,
+                    global_pose=optimed_orient_mat,
                     jaw_pose=tensor2variable(data["jaw_pose"], device),
                     left_hand_pose=tensor2variable(
                         data["left_hand_pose"], device),
@@ -316,8 +321,8 @@ def generate_model(in_path, model_type):
             f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.glb")
 
         smpl_info = {'betas': optimed_betas,
-                     'pose': optimed_pose,
-                     'orient': optimed_orient,
+                     'pose': optimed_pose_mat,
+                     'orient': optimed_orient_mat,
                      'trans': optimed_trans}
 
         np.save(

lib/common/train_util.py

@@ -32,6 +32,8 @@ import os
 from termcolor import colored
 
 
+
+
 def reshape_sample_tensor(sample_tensor, num_views):
     if num_views == 1:
         return sample_tensor

lib/dataset/TestDataset.py

@@ -240,6 +240,11 @@ class TestDataset():
         # body_pose - [1, 23, 3, 3] / [1, 21, 3, 3]
         # global_orient - [1, 1, 3, 3]
         # smpl_verts - [1, 6890, 3] / [1, 10475, 3]
+        
+        # from rot_mat to rot_6d for better optimization
+        N_body = data_dict["body_pose"].shape[1]
+        data_dict["body_pose"] =  data_dict["body_pose"][:, :, :, :2].reshape(1, N_body,-1)
+        data_dict["global_orient"] =  data_dict["global_orient"][:, :, :, :2].reshape(1, 1,-1)
 
         return data_dict
 

lib/dataset/mesh_util.py

@@ -44,6 +44,22 @@ from pytorch3d.loss import (
 
 from huggingface_hub import hf_hub_download, hf_hub_url, cached_download
 
+def rot6d_to_rotmat(x):
+    """Convert 6D rotation representation to 3x3 rotation matrix.
+    Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
+    Input:
+        (B,6) Batch of 6-D rotation representations
+    Output:
+        (B,3,3) Batch of corresponding rotation matrices
+    """
+    x = x.view(-1, 3, 2)
+    a1 = x[:, :, 0]
+    a2 = x[:, :, 1]
+    b1 = F.normalize(a1)
+    b2 = F.normalize(a2 - torch.einsum("bi,bi->b", b1, a2).unsqueeze(-1) * b1)
+    b3 = torch.cross(b1, b2)
+    return torch.stack((b1, b2, b3), dim=-1)
+
 
 def tensor2variable(tensor, device):
     # [1,23,3,3]