Merge branch 'main' into main

mmhuman3d/models/architectures/builder.py

@@ -28,6 +28,7 @@ def build_from_cfg(cfg, registry, default_args=None):
     name='SMPLXImageBodyModelEstimator', module=SMPLXImageBodyModelEstimator)
 ARCHITECTURES.register_module(
     name='CliffImageBodyModelEstimator', module=CliffImageBodyModelEstimator)
+ARCHITECTURES.register_module(name='PyMAFX', module=PyMAFX)
 
 
 def build_architecture(cfg):

mmhuman3d/models/heads/builder.py

@@ -18,6 +18,8 @@
 HEADS.register_module(name='ExPoseHandHead', module=ExPoseHandHead)
 HEADS.register_module(name='ExPoseFaceHead', module=ExPoseFaceHead)
 HEADS.register_module(name='CliffHead', module=CliffHead)
+HEADS.register_module(name='PyMAFXHead', module=PyMAFXHead)
+HEADS.register_module(name='Regressor', module=Regressor)
 
 
 def build_head(cfg):

mmhuman3d/utils/geometry.py

@@ -427,3 +427,68 @@ def cam_crop2full(crop_cam, center, scale, full_img_shape, focal_length):
     ty = (2 * (cy - img_h / 2.) / bs) + crop_cam[:, 2]
     full_cam = torch.stack([tx, ty, tz], dim=-1)
     return full_cam
+
+
+def projection(pred_joints, pred_camera, iwp_mode=True):
+    """Project 3D points on the image plane based on the given camera info,
+    Identity rotation and Weak Perspective (IWP) camera is used when
+    iwp_mode = True
+    """
+    batch_size = pred_joints.shape[0]
+    if iwp_mode:
+        cam_sxy = pred_camera['cam_sxy']
+        pred_cam_t = torch.stack([
+            cam_sxy[:, 1], cam_sxy[:, 2], 2 * 5000. /
+            (224. * cam_sxy[:, 0] + 1e-9)
+        ],
+                                 dim=-1)
+
+        camera_center = torch.zeros(batch_size, 2)
+        pred_keypoints_2d = perspective_projection(
+            pred_joints,
+            rotation=torch.eye(3).unsqueeze(0).expand(batch_size, -1, -1).to(
+                pred_joints.device),
+            translation=pred_cam_t,
+            focal_length=5000.,
+            camera_center=camera_center)
+
+    else:
+        raise NotImplementedError
+    return pred_keypoints_2d
+
+
+def compute_twist_rotation(rotation_matrix, twist_axis):
+    '''
+    Compute the twist component of given rotation and twist axis
+    https://stackoverflow.com/questions/3684269/component-of-a-quaternion-rotation-around-an-axis
+    Parameters
+    ----------
+    rotation_matrix : Tensor (B, 3, 3,)
+        The rotation to convert
+    twist_axis : Tensor (B, 3,)
+        The twist axis
+    Returns
+    -------
+    Tensor (B, 3, 3)
+        The twist rotation
+    '''
+    quaternion = rotation_matrix_to_quaternion(rotation_matrix)
+
+    twist_axis = twist_axis / (
+        torch.norm(twist_axis, dim=1, keepdim=True) + 1e-9)
+
+    projection = torch.einsum('bi,bi->b', twist_axis,
+                              quaternion[:, 1:]).unsqueeze(-1) * twist_axis
+
+    twist_quaternion = torch.cat([quaternion[:, 0:1], projection], dim=1)
+    twist_quaternion = twist_quaternion / (
+        torch.norm(twist_quaternion, dim=1, keepdim=True) + 1e-9)
+
+    twist_rotation = quat_to_rotmat(twist_quaternion)
+    twist_aa = quaternion_to_angle_axis(twist_quaternion)
+
+    twist_angle = torch.sum(
+        twist_aa, dim=1, keepdim=True) / torch.sum(
+            twist_axis, dim=1, keepdim=True)
+
+    return twist_rotation, twist_angle