add huggingface implementation

README.md

@@ -4,7 +4,8 @@ emoji: ⚡
 colorFrom: blue
 colorTo: indigo
 sdk: gradio
-sdk_version: 4.19.2
+python_version: 3.8
+sdk_version: 4.16.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,130 @@
+import os
+import sys
+import os.path as osp
+from pathlib import Path
+import cv2
+import gradio as gr
+import torch
+import math
+
+try:
+    import mmpose
+except:
+    os.system('pip install /home/user/app/main/transformer_utils')
+
+os.system('cp -rf /home/user/app/assets/conversions.py /home/user/.pyenv/versions/3.8.18/lib/python3.8/site-packages/torchgeometry/core/conversions.py')
+DEFAULT_MODEL='smpler_x_h32'
+OUT_FOLDER = '/home/user/app/demo_out'
+os.makedirs(OUT_FOLDER, exist_ok=True)
+num_gpus = 1 if torch.cuda.is_available() else -1
+print("!!!", torch.cuda.is_available())      
+print(torch.cuda.device_count())   
+print(torch.version.cuda)  
+index = torch.cuda.current_device()
+print(index)  
+print(torch.cuda.get_device_name(index))
+from main.inference import Inferer
+inferer = Inferer(DEFAULT_MODEL, num_gpus, OUT_FOLDER)
+
+def infer(video_input, in_threshold=0.5, num_people="Single person", render_mesh=False):
+    os.system(f'rm -rf {OUT_FOLDER}/*')
+    multi_person = False if (num_people == "Single person") else True
+    cap = cv2.VideoCapture(video_input)
+    fps = math.ceil(cap.get(5))
+    width = int(cap.get(3))
+    height = int(cap.get(4))
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_path = osp.join(OUT_FOLDER, f'out.m4v')
+    final_video_path = osp.join(OUT_FOLDER, f'out.mp4')
+    video_output = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
+    success = 1
+    frame = 0
+    while success:
+        success, original_img = cap.read()
+        if not success:
+            break
+        frame += 1
+        img, mesh_paths, smplx_paths = inferer.infer(original_img, in_threshold, frame, multi_person, not(render_mesh))
+        video_output.write(img)
+    cap.release()
+    video_output.release()
+    cv2.destroyAllWindows()
+    os.system(f'ffmpeg -i {video_path} -c copy {final_video_path}')
+
+    #Compress mesh and smplx files
+    save_path_mesh = os.path.join(OUT_FOLDER, 'mesh')
+    save_mesh_file = os.path.join(OUT_FOLDER, 'mesh.zip')
+    os.makedirs(save_path_mesh, exist_ok= True)
+    save_path_smplx = os.path.join(OUT_FOLDER, 'smplx')
+    save_smplx_file = os.path.join(OUT_FOLDER, 'smplx.zip')
+    os.makedirs(save_path_smplx, exist_ok= True)
+    os.system(f'zip -r {save_mesh_file} {save_path_mesh}')
+    os.system(f'zip -r {save_smplx_file} {save_path_smplx}')
+    return video_path, save_mesh_file, save_smplx_file
+
+TITLE = '''<h1 align="center">SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation</h1>'''
+VIDEO = '''
+<center><iframe width="960" height="540" 
+src="https://www.youtube.com/embed/DepTqbPpVzY?si=qSeQuX-bgm_rON7E"title="SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen>
+</iframe>
+</center><br>'''
+DESCRIPTION = '''
+<b>Official Gradio demo</b> for <a href="https://caizhongang.com/projects/SMPLer-X/"><b>SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation</b></a>.<br>
+<p>
+Note: You can drop a video at the panel (or select one of the examples) 
+    then you will get the 3D reconstructions of the detected human. ).
+</p>
+'''
+
+with gr.Blocks(title="SMPLer-X", css=".gradio-container") as demo:
+
+    gr.Markdown(TITLE)
+    gr.HTML(VIDEO)
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.Video(label="Input video", elem_classes="video")
+            threshold = gr.Slider(0, 1.0, value=0.5, label='BBox detection threshold')
+        with gr.Column(scale=2):
+            num_people = gr.Radio(
+                choices=["Single person", "Multiple people"],
+                value="Single person",
+                label="Number of people",
+                info="Choose how many people are there in the video. Choose 'single person' for faster inference.",
+                interactive=True,
+                scale=1,)
+            gr.HTML("""<br/>""")
+            mesh_as_vertices = gr.Checkbox(
+                label="Render mesh", 
+                info="By default, the reconstructions of human bodies are expressed as vertices for faster inference. Check this option if you want to render the human body with mesh.",
+                interactive=True,
+                scale=1,)
+
+            send_button = gr.Button("Infer")
+    gr.HTML("""<br/>""")
+
+    with gr.Row():
+        with gr.Column():
+            video_output = gr.Video(elem_classes="video")
+        with gr.Column():
+            meshes_output = gr.File(label="3D meshes")
+            smplx_output = gr.File(label= "SMPL-X models")
+    # example_images = gr.Examples([])
+    send_button.click(fn=infer, inputs=[video_input, threshold, num_people, mesh_as_vertices], outputs=[video_output, meshes_output, smplx_output])
+    # with gr.Row():
+    example_videos = gr.Examples([
+        ['/home/user/app/assets/01.mp4'], 
+        ['/home/user/app/assets/02.mp4'], 
+        ['/home/user/app/assets/03.mp4'],
+        ['/home/user/app/assets/04.mp4'],
+        ['/home/user/app/assets/05.mp4'], 
+        ['/home/user/app/assets/06.mp4'], 
+        ['/home/user/app/assets/07.mp4'],
+        ['/home/user/app/assets/08.mp4'],
+        ['/home/user/app/assets/09.mp4'],
+        ], 
+        inputs=[video_input, 0.5])
+
+#demo.queue()
+demo.launch(debug=True)

assets/conversions.py

@@ -0,0 +1,523 @@
+import torch
+import torch.nn as nn
+
+import torchgeometry as tgm
+
+__all__ = [
+    # functional api
+    "pi",
+    "rad2deg",
+    "deg2rad",
+    "convert_points_from_homogeneous",
+    "convert_points_to_homogeneous",
+    "angle_axis_to_rotation_matrix",
+    "rotation_matrix_to_angle_axis",
+    "rotation_matrix_to_quaternion",
+    "quaternion_to_angle_axis",
+    "angle_axis_to_quaternion",
+    "rtvec_to_pose",
+    # layer api
+    "RadToDeg",
+    "DegToRad",
+    "ConvertPointsFromHomogeneous",
+    "ConvertPointsToHomogeneous",
+]
+
+
+"""Constant with number pi
+"""
+pi = torch.Tensor([3.14159265358979323846])
+
+
+def rad2deg(tensor):
+    r"""Function that converts angles from radians to degrees.
+
+    See :class:`~torchgeometry.RadToDeg` for details.
+
+    Args:
+        tensor (Tensor): Tensor of arbitrary shape.
+
+    Returns:
+        Tensor: Tensor with same shape as input.
+
+    Example:
+        >>> input = tgm.pi * torch.rand(1, 3, 3)
+        >>> output = tgm.rad2deg(input)
+    """
+    if not torch.is_tensor(tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}"
+                        .format(type(tensor)))
+
+    return 180. * tensor / pi.to(tensor.device).type(tensor.dtype)
+
+
+def deg2rad(tensor):
+    r"""Function that converts angles from degrees to radians.
+
+    See :class:`~torchgeometry.DegToRad` for details.
+
+    Args:
+        tensor (Tensor): Tensor of arbitrary shape.
+
+    Returns:
+        Tensor: Tensor with same shape as input.
+
+    Examples::
+
+        >>> input = 360. * torch.rand(1, 3, 3)
+        >>> output = tgm.deg2rad(input)
+    """
+    if not torch.is_tensor(tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}"
+                        .format(type(tensor)))
+
+    return tensor * pi.to(tensor.device).type(tensor.dtype) / 180.
+
+
+def convert_points_from_homogeneous(points):
+    r"""Function that converts points from homogeneous to Euclidean space.
+
+    See :class:`~torchgeometry.ConvertPointsFromHomogeneous` for details.
+
+    Examples::
+
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> output = tgm.convert_points_from_homogeneous(input)  # BxNx2
+    """
+    if not torch.is_tensor(points):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(points)))
+    if len(points.shape) < 2:
+        raise ValueError("Input must be at least a 2D tensor. Got {}".format(
+            points.shape))
+
+    return points[..., :-1] / points[..., -1:]
+
+
+def convert_points_to_homogeneous(points):
+    r"""Function that converts points from Euclidean to homogeneous space.
+
+    See :class:`~torchgeometry.ConvertPointsToHomogeneous` for details.
+
+    Examples::
+
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> output = tgm.convert_points_to_homogeneous(input)  # BxNx4
+    """
+    if not torch.is_tensor(points):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(points)))
+    if len(points.shape) < 2:
+        raise ValueError("Input must be at least a 2D tensor. Got {}".format(
+            points.shape))
+
+    return nn.functional.pad(points, (0, 1), "constant", 1.0)
+
+
+def angle_axis_to_rotation_matrix(angle_axis):
+    """Convert 3d vector of axis-angle rotation to 4x4 rotation matrix
+
+    Args:
+        angle_axis (Tensor): tensor of 3d vector of axis-angle rotations.
+
+    Returns:
+        Tensor: tensor of 4x4 rotation matrices.
+
+    Shape:
+        - Input: :math:`(N, 3)`
+        - Output: :math:`(N, 4, 4)`
+
+    Example:
+        >>> input = torch.rand(1, 3)  # Nx3
+        >>> output = tgm.angle_axis_to_rotation_matrix(input)  # Nx4x4
+    """
+    def _compute_rotation_matrix(angle_axis, theta2, eps=1e-6):
+        # We want to be careful to only evaluate the square root if the
+        # norm of the angle_axis vector is greater than zero. Otherwise
+        # we get a division by zero.
+        k_one = 1.0
+        theta = torch.sqrt(theta2)
+        wxyz = angle_axis / (theta + eps)
+        wx, wy, wz = torch.chunk(wxyz, 3, dim=1)
+        cos_theta = torch.cos(theta)
+        sin_theta = torch.sin(theta)
+
+        r00 = cos_theta + wx * wx * (k_one - cos_theta)
+        r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
+        r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
+        r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
+        r11 = cos_theta + wy * wy * (k_one - cos_theta)
+        r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
+        r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
+        r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
+        r22 = cos_theta + wz * wz * (k_one - cos_theta)
+        rotation_matrix = torch.cat(
+            [r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
+        return rotation_matrix.view(-1, 3, 3)
+
+    def _compute_rotation_matrix_taylor(angle_axis):
+        rx, ry, rz = torch.chunk(angle_axis, 3, dim=1)
+        k_one = torch.ones_like(rx)
+        rotation_matrix = torch.cat(
+            [k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1)
+        return rotation_matrix.view(-1, 3, 3)
+
+    # stolen from ceres/rotation.h
+
+    _angle_axis = torch.unsqueeze(angle_axis, dim=1)
+    theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
+    theta2 = torch.squeeze(theta2, dim=1)
+
+    # compute rotation matrices
+    rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
+    rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)
+
+    # create mask to handle both cases
+    eps = 1e-6
+    mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device)
+    mask_pos = (mask).type_as(theta2)
+    mask_neg = (mask == False).type_as(theta2)  # noqa
+
+    # create output pose matrix
+    batch_size = angle_axis.shape[0]
+    rotation_matrix = torch.eye(4).to(angle_axis.device).type_as(angle_axis)
+    rotation_matrix = rotation_matrix.view(1, 4, 4).repeat(batch_size, 1, 1)
+    # fill output matrix with masked values
+    rotation_matrix[..., :3, :3] = \
+        mask_pos * rotation_matrix_normal + mask_neg * rotation_matrix_taylor
+    return rotation_matrix  # Nx4x4
+
+
+def rtvec_to_pose(rtvec):
+    """
+    Convert axis-angle rotation and translation vector to 4x4 pose matrix
+
+    Args:
+        rtvec (Tensor): Rodrigues vector transformations
+
+    Returns:
+        Tensor: transformation matrices
+
+    Shape:
+        - Input: :math:`(N, 6)`
+        - Output: :math:`(N, 4, 4)`
+
+    Example:
+        >>> input = torch.rand(3, 6)  # Nx6
+        >>> output = tgm.rtvec_to_pose(input)  # Nx4x4
+    """
+    assert rtvec.shape[-1] == 6, 'rtvec=[rx, ry, rz, tx, ty, tz]'
+    pose = angle_axis_to_rotation_matrix(rtvec[..., :3])
+    pose[..., :3, 3] = rtvec[..., 3:]
+    return pose
+
+
+def rotation_matrix_to_angle_axis(rotation_matrix):
+    """Convert 3x4 rotation matrix to Rodrigues vector
+
+    Args:
+        rotation_matrix (Tensor): rotation matrix.
+
+    Returns:
+        Tensor: Rodrigues vector transformation.
+
+    Shape:
+        - Input: :math:`(N, 3, 4)`
+        - Output: :math:`(N, 3)`
+
+    Example:
+        >>> input = torch.rand(2, 3, 4)  # Nx4x4
+        >>> output = tgm.rotation_matrix_to_angle_axis(input)  # Nx3
+    """
+    # todo add check that matrix is a valid rotation matrix
+    quaternion = rotation_matrix_to_quaternion(rotation_matrix)
+    return quaternion_to_angle_axis(quaternion)
+
+
+def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
+    """Convert 3x4 rotation matrix to 4d quaternion vector
+
+    This algorithm is based on algorithm described in
+    https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201
+
+    Args:
+        rotation_matrix (Tensor): the rotation matrix to convert.
+
+    Return:
+        Tensor: the rotation in quaternion
+
+    Shape:
+        - Input: :math:`(N, 3, 4)`
+        - Output: :math:`(N, 4)`
+
+    Example:
+        >>> input = torch.rand(4, 3, 4)  # Nx3x4
+        >>> output = tgm.rotation_matrix_to_quaternion(input)  # Nx4
+    """
+    if not torch.is_tensor(rotation_matrix):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(rotation_matrix)))
+
+    if len(rotation_matrix.shape) > 3:
+        raise ValueError(
+            "Input size must be a three dimensional tensor. Got {}".format(
+                rotation_matrix.shape))
+    if not rotation_matrix.shape[-2:] == (3, 4):
+        raise ValueError(
+            "Input size must be a N x 3 x 4  tensor. Got {}".format(
+                rotation_matrix.shape))
+
+    rmat_t = torch.transpose(rotation_matrix, 1, 2)
+
+    mask_d2 = rmat_t[:, 2, 2] < eps
+
+    mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
+    mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]
+
+    t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
+    q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
+                      t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
+                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1)
+    t0_rep = t0.repeat(4, 1).t()
+
+    t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
+    q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
+                      rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
+                      t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1)
+    t1_rep = t1.repeat(4, 1).t()
+
+    t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
+    q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
+                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
+                      rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1)
+    t2_rep = t2.repeat(4, 1).t()
+
+    t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
+    q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
+                      rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
+                      rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1)
+    t3_rep = t3.repeat(4, 1).t()
+
+    mask_c0 = mask_d2 * mask_d0_d1
+    mask_c1 = mask_d2 * ~(mask_d0_d1)
+    mask_c2 = ~(mask_d2) * mask_d0_nd1
+    mask_c3 = ~(mask_d2) * ~(mask_d0_nd1)
+    mask_c0 = mask_c0.view(-1, 1).type_as(q0)
+    mask_c1 = mask_c1.view(-1, 1).type_as(q1)
+    mask_c2 = mask_c2.view(-1, 1).type_as(q2)
+    mask_c3 = mask_c3.view(-1, 1).type_as(q3)
+
+    q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
+    q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 +  # noqa
+                    t2_rep * mask_c2 + t3_rep * mask_c3)  # noqa
+    q *= 0.5
+    return q
+
+
+def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:
+    """Convert quaternion vector to angle axis of rotation.
+
+    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
+
+    Args:
+        quaternion (torch.Tensor): tensor with quaternions.
+
+    Return:
+        torch.Tensor: tensor with angle axis of rotation.
+
+    Shape:
+        - Input: :math:`(*, 4)` where `*` means, any number of dimensions
+        - Output: :math:`(*, 3)`
+
+    Example:
+        >>> quaternion = torch.rand(2, 4)  # Nx4
+        >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion)  # Nx3
+    """
+    if not torch.is_tensor(quaternion):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(quaternion)))
+
+    if not quaternion.shape[-1] == 4:
+        raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}"
+                         .format(quaternion.shape))
+    # unpack input and compute conversion
+    q1: torch.Tensor = quaternion[..., 1]
+    q2: torch.Tensor = quaternion[..., 2]
+    q3: torch.Tensor = quaternion[..., 3]
+    sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3
+
+    sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)
+    cos_theta: torch.Tensor = quaternion[..., 0]
+    two_theta: torch.Tensor = 2.0 * torch.where(
+        cos_theta < 0.0,
+        torch.atan2(-sin_theta, -cos_theta),
+        torch.atan2(sin_theta, cos_theta))
+
+    k_pos: torch.Tensor = two_theta / sin_theta
+    k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)
+    k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)
+
+    angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]
+    angle_axis[..., 0] += q1 * k
+    angle_axis[..., 1] += q2 * k
+    angle_axis[..., 2] += q3 * k
+    return angle_axis
+
+# based on:
+# https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py#L138
+
+
+def angle_axis_to_quaternion(angle_axis: torch.Tensor) -> torch.Tensor:
+    """Convert an angle axis to a quaternion.
+
+    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
+
+    Args:
+        angle_axis (torch.Tensor): tensor with angle axis.
+
+    Return:
+        torch.Tensor: tensor with quaternion.
+
+    Shape:
+        - Input: :math:`(*, 3)` where `*` means, any number of dimensions
+        - Output: :math:`(*, 4)`
+
+    Example:
+        >>> angle_axis = torch.rand(2, 4)  # Nx4
+        >>> quaternion = tgm.angle_axis_to_quaternion(angle_axis)  # Nx3
+    """
+    if not torch.is_tensor(angle_axis):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(angle_axis)))
+
+    if not angle_axis.shape[-1] == 3:
+        raise ValueError("Input must be a tensor of shape Nx3 or 3. Got {}"
+                         .format(angle_axis.shape))
+    # unpack input and compute conversion
+    a0: torch.Tensor = angle_axis[..., 0:1]
+    a1: torch.Tensor = angle_axis[..., 1:2]
+    a2: torch.Tensor = angle_axis[..., 2:3]
+    theta_squared: torch.Tensor = a0 * a0 + a1 * a1 + a2 * a2
+
+    theta: torch.Tensor = torch.sqrt(theta_squared)
+    half_theta: torch.Tensor = theta * 0.5
+
+    mask: torch.Tensor = theta_squared > 0.0
+    ones: torch.Tensor = torch.ones_like(half_theta)
+
+    k_neg: torch.Tensor = 0.5 * ones
+    k_pos: torch.Tensor = torch.sin(half_theta) / theta
+    k: torch.Tensor = torch.where(mask, k_pos, k_neg)
+    w: torch.Tensor = torch.where(mask, torch.cos(half_theta), ones)
+
+    quaternion: torch.Tensor = torch.zeros_like(angle_axis)
+    quaternion[..., 0:1] += a0 * k
+    quaternion[..., 1:2] += a1 * k
+    quaternion[..., 2:3] += a2 * k
+    return torch.cat([w, quaternion], dim=-1)
+
+# TODO: add below funtionalities
+#  - pose_to_rtvec
+
+
+# layer api
+
+
+class RadToDeg(nn.Module):
+    r"""Creates an object that converts angles from radians to degrees.
+
+    Args:
+        tensor (Tensor): Tensor of arbitrary shape.
+
+    Returns:
+        Tensor: Tensor with same shape as input.
+
+    Examples::
+
+        >>> input = tgm.pi * torch.rand(1, 3, 3)
+        >>> output = tgm.RadToDeg()(input)
+    """
+
+    def __init__(self):
+        super(RadToDeg, self).__init__()
+
+    def forward(self, input):
+        return rad2deg(input)
+
+
+class DegToRad(nn.Module):
+    r"""Function that converts angles from degrees to radians.
+
+    Args:
+        tensor (Tensor): Tensor of arbitrary shape.
+
+    Returns:
+        Tensor: Tensor with same shape as input.
+
+    Examples::
+
+        >>> input = 360. * torch.rand(1, 3, 3)
+        >>> output = tgm.DegToRad()(input)
+    """
+
+    def __init__(self):
+        super(DegToRad, self).__init__()
+
+    def forward(self, input):
+        return deg2rad(input)
+
+
+class ConvertPointsFromHomogeneous(nn.Module):
+    r"""Creates a transformation that converts points from homogeneous to
+    Euclidean space.
+
+    Args:
+        points (Tensor): tensor of N-dimensional points.
+
+    Returns:
+        Tensor: tensor of N-1-dimensional points.
+
+    Shape:
+        - Input: :math:`(B, D, N)` or :math:`(D, N)`
+        - Output: :math:`(B, D, N + 1)` or :math:`(D, N + 1)`
+
+    Examples::
+
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> transform = tgm.ConvertPointsFromHomogeneous()
+        >>> output = transform(input)  # BxNx2
+    """
+
+    def __init__(self):
+        super(ConvertPointsFromHomogeneous, self).__init__()
+
+    def forward(self, input):
+        return convert_points_from_homogeneous(input)
+
+
+class ConvertPointsToHomogeneous(nn.Module):
+    r"""Creates a transformation to convert points from Euclidean to
+    homogeneous space.
+
+    Args:
+        points (Tensor): tensor of N-dimensional points.
+
+    Returns:
+        Tensor: tensor of N+1-dimensional points.
+
+    Shape:
+        - Input: :math:`(B, D, N)` or :math:`(D, N)`
+        - Output: :math:`(B, D, N + 1)` or :math:`(D, N + 1)`
+
+    Examples::
+
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> transform = tgm.ConvertPointsToHomogeneous()
+        >>> output = transform(input)  # BxNx4
+    """
+
+    def __init__(self):
+        super(ConvertPointsToHomogeneous, self).__init__()
+
+    def forward(self, input):
+        return convert_points_to_homogeneous(input)

common/base.py

@@ -0,0 +1,86 @@
+import os.path as osp
+import math
+import abc
+from torch.utils.data import DataLoader
+import torch.optim
+import torchvision.transforms as transforms
+from timer import Timer
+from logger import colorlogger
+from torch.nn.parallel.data_parallel import DataParallel
+from config import cfg
+from SMPLer_X import get_model
+
+# ddp
+import torch.distributed as dist
+from torch.utils.data import DistributedSampler
+import torch.utils.data.distributed
+from utils.distribute_utils import (
+    get_rank, is_main_process, time_synchronized, get_group_idx, get_process_groups
+)
+from mmcv.runner import get_dist_info
+
+class Base(object):
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, log_name='logs.txt'):
+        self.cur_epoch = 0
+
+        # timer
+        self.tot_timer = Timer()
+        self.gpu_timer = Timer()
+        self.read_timer = Timer()
+
+        # logger
+        self.logger = colorlogger(cfg.log_dir, log_name=log_name)
+
+    @abc.abstractmethod
+    def _make_batch_generator(self):
+        return
+
+    @abc.abstractmethod
+    def _make_model(self):
+        return
+
+class Demoer(Base):
+    def __init__(self, test_epoch=None):
+        if test_epoch is not None:
+            self.test_epoch = int(test_epoch)
+        super(Demoer, self).__init__(log_name='test_logs.txt')
+
+    def _make_batch_generator(self, demo_scene):
+        # data load and construct batch generator
+        self.logger.info("Creating dataset...")
+        from data.UBody.UBody import UBody
+        testset_loader = UBody(transforms.ToTensor(), "demo", demo_scene) # eval(demoset)(transforms.ToTensor(), "demo")
+        batch_generator = DataLoader(dataset=testset_loader, batch_size=cfg.num_gpus * cfg.test_batch_size,
+                                     shuffle=False, num_workers=cfg.num_thread, pin_memory=True)
+
+        self.testset = testset_loader
+        self.batch_generator = batch_generator
+
+    def _make_model(self):
+        self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path))
+
+        # prepare network
+        self.logger.info("Creating graph...")
+        model = get_model('test')
+        model = DataParallel(model).to(cfg.device)
+        ckpt = torch.load(cfg.pretrained_model_path, map_location=cfg.device)
+
+        from collections import OrderedDict
+        new_state_dict = OrderedDict()
+        for k, v in ckpt['network'].items():
+            if 'module' not in k:
+                k = 'module.' + k
+            k = k.replace('module.backbone', 'module.encoder').replace('body_rotation_net', 'body_regressor').replace(
+                'hand_rotation_net', 'hand_regressor')
+            new_state_dict[k] = v
+        model.load_state_dict(new_state_dict, strict=False)
+        model.eval()
+
+        self.model = model
+
+    def _evaluate(self, outs, cur_sample_idx):
+        eval_result = self.testset.evaluate(outs, cur_sample_idx)
+        return eval_result
+

common/logger.py

@@ -0,0 +1,50 @@
+import logging
+import os
+
+OK = '\033[92m'
+WARNING = '\033[93m'
+FAIL = '\033[91m'
+END = '\033[0m'
+
+PINK = '\033[95m'
+BLUE = '\033[94m'
+GREEN = OK
+RED = FAIL
+WHITE = END
+YELLOW = WARNING
+
+class colorlogger():
+    def __init__(self, log_dir, log_name='train_logs.txt'):
+        # set log
+        self._logger = logging.getLogger(log_name)
+        self._logger.setLevel(logging.INFO)
+        log_file = os.path.join(log_dir, log_name)
+        if not os.path.exists(log_dir):
+            os.makedirs(log_dir)
+        file_log = logging.FileHandler(log_file, mode='a')
+        file_log.setLevel(logging.INFO)
+        console_log = logging.StreamHandler()
+        console_log.setLevel(logging.INFO)
+        formatter = logging.Formatter(
+            "{}%(asctime)s{} %(message)s".format(GREEN, END),
+            "%m-%d %H:%M:%S")
+        file_log.setFormatter(formatter)
+        console_log.setFormatter(formatter)
+        self._logger.addHandler(file_log)
+        self._logger.addHandler(console_log)
+
+    def debug(self, msg):
+        self._logger.debug(str(msg))
+
+    def info(self, msg):
+        self._logger.info(str(msg))
+
+    def warning(self, msg):
+        self._logger.warning(WARNING + 'WRN: ' + str(msg) + END)
+
+    def critical(self, msg):
+        self._logger.critical(RED + 'CRI: ' + str(msg) + END)
+
+    def error(self, msg):
+        self._logger.error(RED + 'ERR: ' + str(msg) + END)
+

common/nets/layer.py

@@ -0,0 +1,53 @@
+import torch.nn as nn
+
+def make_linear_layers(feat_dims, relu_final=True, use_bn=False):
+    layers = []
+    for i in range(len(feat_dims)-1):
+        layers.append(nn.Linear(feat_dims[i], feat_dims[i+1]))
+
+        # Do not use ReLU for final estimation
+        if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and relu_final):
+            if use_bn:
+                layers.append(nn.BatchNorm1d(feat_dims[i+1]))
+            layers.append(nn.ReLU(inplace=True))
+
+    return nn.Sequential(*layers)
+
+def make_conv_layers(feat_dims, kernel=3, stride=1, padding=1, bnrelu_final=True):
+    layers = []
+    for i in range(len(feat_dims)-1):
+        layers.append(
+            nn.Conv2d(
+                in_channels=feat_dims[i],
+                out_channels=feat_dims[i+1],
+                kernel_size=kernel,
+                stride=stride,
+                padding=padding
+                ))
+        # Do not use BN and ReLU for final estimation
+        if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final):
+            layers.append(nn.BatchNorm2d(feat_dims[i+1]))
+            layers.append(nn.ReLU(inplace=True))
+
+    return nn.Sequential(*layers)
+
+def make_deconv_layers(feat_dims, bnrelu_final=True):
+    layers = []
+    for i in range(len(feat_dims)-1):
+        layers.append(
+            nn.ConvTranspose2d(
+                in_channels=feat_dims[i],
+                out_channels=feat_dims[i+1],
+                kernel_size=4,
+                stride=2,
+                padding=1,
+                output_padding=0,
+                bias=False))
+
+        # Do not use BN and ReLU for final estimation
+        if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final):
+            layers.append(nn.BatchNorm2d(feat_dims[i+1]))
+            layers.append(nn.ReLU(inplace=True))
+
+    return nn.Sequential(*layers)
+

common/nets/loss.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+
+class CoordLoss(nn.Module):
+    def __init__(self):
+        super(CoordLoss, self).__init__()
+
+    def forward(self, coord_out, coord_gt, valid, is_3D=None):
+        loss = torch.abs(coord_out - coord_gt) * valid
+        if is_3D is not None:
+            loss_z = loss[:,:,2:] * is_3D[:,None,None].float()
+            loss = torch.cat((loss[:,:,:2], loss_z),2)
+        return loss
+
+class ParamLoss(nn.Module):
+    def __init__(self):
+        super(ParamLoss, self).__init__()
+
+    def forward(self, param_out, param_gt, valid):
+        loss = torch.abs(param_out - param_gt) * valid
+        return loss
+
+class CELoss(nn.Module):
+    def __init__(self):
+        super(CELoss, self).__init__()
+        self.ce_loss = nn.CrossEntropyLoss(reduction='none')
+
+    def forward(self, out, gt_index):
+        loss = self.ce_loss(out, gt_index)
+        return loss

common/nets/smpler_x.py

@@ -0,0 +1,172 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from nets.layer import make_conv_layers, make_linear_layers, make_deconv_layers
+from utils.transforms import sample_joint_features, soft_argmax_2d, soft_argmax_3d
+from utils.human_models import smpl_x
+from config import cfg
+from mmcv.ops.roi_align import roi_align
+
+class PositionNet(nn.Module):
+    def __init__(self, part, feat_dim=768):
+        super(PositionNet, self).__init__()
+        if part == 'body':
+            self.joint_num = len(smpl_x.pos_joint_part['body'])
+            self.hm_shape = cfg.output_hm_shape
+        elif part == 'hand':
+            self.joint_num = len(smpl_x.pos_joint_part['rhand'])
+            self.hm_shape = cfg.output_hand_hm_shape
+        self.conv = make_conv_layers([feat_dim, self.joint_num * self.hm_shape[0]], kernel=1, stride=1, padding=0, bnrelu_final=False)
+
+    def forward(self, img_feat):
+        joint_hm = self.conv(img_feat).view(-1, self.joint_num, self.hm_shape[0], self.hm_shape[1], self.hm_shape[2])
+        joint_coord = soft_argmax_3d(joint_hm)
+        joint_hm = F.softmax(joint_hm.view(-1, self.joint_num, self.hm_shape[0] * self.hm_shape[1] * self.hm_shape[2]), 2)
+        joint_hm = joint_hm.view(-1, self.joint_num, self.hm_shape[0], self.hm_shape[1], self.hm_shape[2])
+        return joint_hm, joint_coord
+
+class HandRotationNet(nn.Module):
+    def __init__(self, part, feat_dim = 768):
+        super(HandRotationNet, self).__init__()
+        self.part = part
+        self.joint_num = len(smpl_x.pos_joint_part['rhand'])
+        self.hand_conv = make_conv_layers([feat_dim, 512], kernel=1, stride=1, padding=0)
+        self.hand_pose_out = make_linear_layers([self.joint_num * 515, len(smpl_x.orig_joint_part['rhand']) * 6], relu_final=False)
+        self.feat_dim = feat_dim
+
+    def forward(self, img_feat, joint_coord_img):
+        batch_size = img_feat.shape[0]
+        img_feat = self.hand_conv(img_feat)
+        img_feat_joints = sample_joint_features(img_feat, joint_coord_img[:, :, :2])
+        feat = torch.cat((img_feat_joints, joint_coord_img), 2)  # batch_size, joint_num, 512+3
+        hand_pose = self.hand_pose_out(feat.view(batch_size, -1))
+        return hand_pose
+
+class BodyRotationNet(nn.Module):
+    def __init__(self, feat_dim = 768):
+        super(BodyRotationNet, self).__init__()
+        self.joint_num = len(smpl_x.pos_joint_part['body'])
+        self.body_conv = make_linear_layers([feat_dim, 512], relu_final=False)
+        self.root_pose_out = make_linear_layers([self.joint_num * (512+3), 6], relu_final=False)
+        self.body_pose_out = make_linear_layers(
+            [self.joint_num * (512+3), (len(smpl_x.orig_joint_part['body']) - 1) * 6], relu_final=False)  # without root
+        self.shape_out = make_linear_layers([feat_dim, smpl_x.shape_param_dim], relu_final=False)
+        self.cam_out = make_linear_layers([feat_dim, 3], relu_final=False)
+        self.feat_dim = feat_dim
+
+    def forward(self, body_pose_token, shape_token, cam_token, body_joint_img):
+        batch_size = body_pose_token.shape[0]
+
+        # shape parameter
+        shape_param = self.shape_out(shape_token)
+
+        # camera parameter
+        cam_param = self.cam_out(cam_token)
+
+        # body pose parameter
+        body_pose_token = self.body_conv(body_pose_token)
+        body_pose_token = torch.cat((body_pose_token, body_joint_img), 2)
+        root_pose = self.root_pose_out(body_pose_token.view(batch_size, -1))
+        body_pose = self.body_pose_out(body_pose_token.view(batch_size, -1))
+
+        return root_pose, body_pose, shape_param, cam_param
+
+class FaceRegressor(nn.Module):
+    def __init__(self, feat_dim=768):
+        super(FaceRegressor, self).__init__()
+        self.expr_out = make_linear_layers([feat_dim, smpl_x.expr_code_dim], relu_final=False)
+        self.jaw_pose_out = make_linear_layers([feat_dim, 6], relu_final=False)
+
+    def forward(self, expr_token, jaw_pose_token):
+        expr_param = self.expr_out(expr_token)  # expression parameter
+        jaw_pose = self.jaw_pose_out(jaw_pose_token)  # jaw pose parameter
+        return expr_param, jaw_pose
+
+class BoxNet(nn.Module):
+    def __init__(self, feat_dim=768):
+        super(BoxNet, self).__init__()
+        self.joint_num = len(smpl_x.pos_joint_part['body'])
+        self.deconv = make_deconv_layers([feat_dim + self.joint_num * cfg.output_hm_shape[0], 256, 256, 256])
+        self.bbox_center = make_conv_layers([256, 3], kernel=1, stride=1, padding=0, bnrelu_final=False)
+        self.lhand_size = make_linear_layers([256, 256, 2], relu_final=False)
+        self.rhand_size = make_linear_layers([256, 256, 2], relu_final=False)
+        self.face_size = make_linear_layers([256, 256, 2], relu_final=False)
+
+    def forward(self, img_feat, joint_hm):
+        joint_hm = joint_hm.view(joint_hm.shape[0], joint_hm.shape[1] * cfg.output_hm_shape[0], cfg.output_hm_shape[1], cfg.output_hm_shape[2])
+        img_feat = torch.cat((img_feat, joint_hm), 1)
+        img_feat = self.deconv(img_feat)
+
+        # bbox center
+        bbox_center_hm = self.bbox_center(img_feat)
+        bbox_center = soft_argmax_2d(bbox_center_hm)
+        lhand_center, rhand_center, face_center = bbox_center[:, 0, :], bbox_center[:, 1, :], bbox_center[:, 2, :]
+
+        # bbox size
+        lhand_feat = sample_joint_features(img_feat, lhand_center[:, None, :].detach())[:, 0, :]
+        lhand_size = self.lhand_size(lhand_feat)
+        rhand_feat = sample_joint_features(img_feat, rhand_center[:, None, :].detach())[:, 0, :]
+        rhand_size = self.rhand_size(rhand_feat)
+        face_feat = sample_joint_features(img_feat, face_center[:, None, :].detach())[:, 0, :]
+        face_size = self.face_size(face_feat)
+
+        lhand_center = lhand_center / 8
+        rhand_center = rhand_center / 8
+        face_center = face_center / 8
+        return lhand_center, lhand_size, rhand_center, rhand_size, face_center, face_size
+
+class BoxSizeNet(nn.Module):
+    def __init__(self):
+        super(BoxSizeNet, self).__init__()
+        self.lhand_size = make_linear_layers([256, 256, 2], relu_final=False)
+        self.rhand_size = make_linear_layers([256, 256, 2], relu_final=False)
+        self.face_size = make_linear_layers([256, 256, 2], relu_final=False)
+
+    def forward(self, box_fea):
+        # box_fea: [bs, 3, C]
+        lhand_size = self.lhand_size(box_fea[:, 0])
+        rhand_size = self.rhand_size(box_fea[:, 1])
+        face_size = self.face_size(box_fea[:, 2])
+        return lhand_size, rhand_size, face_size
+
+class HandRoI(nn.Module):
+    def __init__(self, feat_dim=768, upscale=4):
+        super(HandRoI, self).__init__()
+        self.upscale = upscale
+        if upscale==1:
+            self.deconv = make_conv_layers([feat_dim, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
+            self.conv = make_conv_layers([feat_dim, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
+        elif upscale==2:
+            self.deconv = make_deconv_layers([feat_dim, feat_dim//2])
+            self.conv = make_conv_layers([feat_dim//2, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
+        elif upscale==4:
+            self.deconv = make_deconv_layers([feat_dim, feat_dim//2, feat_dim//4])
+            self.conv = make_conv_layers([feat_dim//4, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
+        elif upscale==8:
+            self.deconv = make_deconv_layers([feat_dim, feat_dim//2, feat_dim//4, feat_dim//8])
+            self.conv = make_conv_layers([feat_dim//8, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
+
+    def forward(self, img_feat, lhand_bbox, rhand_bbox):
+        lhand_bbox = torch.cat((torch.arange(lhand_bbox.shape[0]).float().to(cfg.device)[:, None], lhand_bbox),
+                               1)  # batch_idx, xmin, ymin, xmax, ymax
+        rhand_bbox = torch.cat((torch.arange(rhand_bbox.shape[0]).float().to(cfg.device)[:, None], rhand_bbox),
+                               1)  # batch_idx, xmin, ymin, xmax, ymax
+        img_feat = self.deconv(img_feat)
+        lhand_bbox_roi = lhand_bbox.clone()
+        lhand_bbox_roi[:, 1] = lhand_bbox_roi[:, 1] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
+        lhand_bbox_roi[:, 2] = lhand_bbox_roi[:, 2] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
+        lhand_bbox_roi[:, 3] = lhand_bbox_roi[:, 3] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
+        lhand_bbox_roi[:, 4] = lhand_bbox_roi[:, 4] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
+        assert (cfg.output_hm_shape[1]*self.upscale, cfg.output_hm_shape[2]*self.upscale) == (img_feat.shape[2], img_feat.shape[3])
+        lhand_img_feat = roi_align(img_feat, lhand_bbox_roi, (cfg.output_hand_hm_shape[1], cfg.output_hand_hm_shape[2]), 1.0, 0, 'avg', False)
+        lhand_img_feat = torch.flip(lhand_img_feat, [3])  # flip to the right hand
+
+        rhand_bbox_roi = rhand_bbox.clone()
+        rhand_bbox_roi[:, 1] = rhand_bbox_roi[:, 1] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
+        rhand_bbox_roi[:, 2] = rhand_bbox_roi[:, 2] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
+        rhand_bbox_roi[:, 3] = rhand_bbox_roi[:, 3] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
+        rhand_bbox_roi[:, 4] = rhand_bbox_roi[:, 4] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
+        rhand_img_feat = roi_align(img_feat, rhand_bbox_roi, (cfg.output_hand_hm_shape[1], cfg.output_hand_hm_shape[2]), 1.0, 0, 'avg', False)
+        hand_img_feat = torch.cat((lhand_img_feat, rhand_img_feat))  # [bs, c, cfg.output_hand_hm_shape[2]*scale, cfg.output_hand_hm_shape[1]*scale]
+        hand_img_feat = self.conv(hand_img_feat)
+        return hand_img_feat

common/timer.py

@@ -0,0 +1,38 @@
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import time
+
+class Timer(object):
+    """A simple timer."""
+    def __init__(self):
+        self.total_time = 0.
+        self.calls = 0
+        self.start_time = 0.
+        self.diff = 0.
+        self.average_time = 0.
+        self.warm_up = 0
+
+    def tic(self):
+        # using time.time instead of time.clock because time time.clock
+        # does not normalize for multithreading
+        self.start_time = time.time()
+
+    def toc(self, average=True):
+        self.diff = time.time() - self.start_time
+        if self.warm_up < 10:
+            self.warm_up += 1
+            return self.diff
+        else:
+            self.total_time += self.diff
+            self.calls += 1
+            self.average_time = self.total_time / self.calls
+
+        if average:
+            return self.average_time
+        else:
+            return self.diff

common/utils/dir.py

@@ -0,0 +1,10 @@
+import os
+import sys
+
+def make_folder(folder_name):
+    os.makedirs(folder_name, exist_ok=True)
+
+def add_pypath(path):
+    if path not in sys.path:
+        sys.path.insert(0, path)
+

common/utils/distribute_utils.py

@@ -0,0 +1,217 @@
+import mmcv
+import os
+import os.path as osp
+import pickle
+import shutil
+import tempfile
+import time
+import torch
+import torch.distributed as dist
+from mmcv.runner import get_dist_info
+import random
+import numpy as np
+import subprocess
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    # torch.set_deterministic(True)
+
+
+def time_synchronized():
+    torch.cuda.synchronize() if torch.cuda.is_available() else None
+    return time.time()
+
+
+def setup_for_distributed(is_master):
+    """This function disables printing when not in master process."""
+    import builtins as __builtin__
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def init_distributed_mode(port = None, master_port=29500):
+    """Initialize slurm distributed training environment.
+
+    If argument ``port`` is not specified, then the master port will be system
+    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
+    environment variable, then a default port ``29500`` will be used.
+
+    Args:
+        backend (str): Backend of torch.distributed.
+        port (int, optional): Master port. Defaults to None.
+    """
+    dist_backend = 'nccl'
+    proc_id = int(os.environ['SLURM_PROCID'])
+    ntasks = int(os.environ['SLURM_NTASKS'])
+    node_list = os.environ['SLURM_NODELIST']
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(proc_id % num_gpus)
+    addr = subprocess.getoutput(
+        f'scontrol show hostname {node_list} | head -n1')
+    # specify master port
+    if port is not None:
+        os.environ['MASTER_PORT'] = str(port)
+    elif 'MASTER_PORT' in os.environ:
+        pass  # use MASTER_PORT in the environment variable
+    else:
+        # 29500 is torch.distributed default port
+        os.environ['MASTER_PORT'] = str(master_port)
+    # use MASTER_ADDR in the environment variable if it already exists
+    if 'MASTER_ADDR' not in os.environ:
+        os.environ['MASTER_ADDR'] = addr
+    os.environ['WORLD_SIZE'] = str(ntasks)
+    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
+    os.environ['RANK'] = str(proc_id)
+    dist.init_process_group(backend=dist_backend)
+
+    distributed = True
+    gpu_idx = proc_id % num_gpus
+
+    return distributed, gpu_idx
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+def get_process_groups():
+    world_size = int(os.environ['WORLD_SIZE'])
+    ranks = list(range(world_size))
+    num_gpus = torch.cuda.device_count()
+    num_nodes = world_size // num_gpus
+    if world_size % num_gpus != 0:
+        raise NotImplementedError('Not implemented for node not fully used.')
+
+    groups = []
+    for node_idx in range(num_nodes):
+        groups.append(ranks[node_idx*num_gpus : (node_idx+1)*num_gpus])
+    process_groups = [torch.distributed.new_group(group) for group in groups]
+
+    return process_groups
+
+def get_group_idx():
+    num_gpus = torch.cuda.device_count()
+    proc_id = get_rank()
+    group_idx = proc_id // num_gpus
+
+    return group_idx
+
+
+def is_main_process():
+    return get_rank() == 0
+
+def cleanup():
+    dist.destroy_process_group()
+
+
+def collect_results(result_part, size, tmpdir=None):
+    rank, world_size = get_dist_info()
+    # create a tmp dir if it is not specified
+    if tmpdir is None:
+        MAX_LEN = 512
+        # 32 is whitespace
+        dir_tensor = torch.full((MAX_LEN, ),
+                                32,
+                                dtype=torch.uint8,
+                                device='cuda')
+        if rank == 0:
+            tmpdir = tempfile.mkdtemp()
+            tmpdir = torch.tensor(
+                bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda')
+            dir_tensor[:len(tmpdir)] = tmpdir
+        dist.broadcast(dir_tensor, 0)
+        tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()
+    else:
+        mmcv.mkdir_or_exist(tmpdir)
+    # dump the part result to the dir
+    mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl'))
+    dist.barrier()
+    # collect all parts
+    if rank != 0:
+        return None
+    else:
+        # load results of all parts from tmp dir
+        part_list = []
+        for i in range(world_size):
+            part_file = osp.join(tmpdir, f'part_{i}.pkl')
+            part_list.append(mmcv.load(part_file))
+        # sort the results
+        ordered_results = []
+        for res in zip(*part_list):
+            ordered_results.extend(list(res))
+        # the dataloader may pad some samples
+        ordered_results = ordered_results[:size]
+        # remove tmp dir
+        shutil.rmtree(tmpdir)
+        return ordered_results
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data:
+            Any picklable object
+    Returns:
+        data_list(list):
+            List of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to('cuda')
+
+    # obtain Tensor size of each rank
+    local_size = torch.tensor([tensor.numel()], device='cuda')
+    size_list = [torch.tensor([0], device='cuda') for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(
+            torch.empty((max_size, ), dtype=torch.uint8, device='cuda'))
+    if local_size != max_size:
+        padding = torch.empty(
+            size=(max_size - local_size, ), dtype=torch.uint8, device='cuda')
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list

common/utils/human_model_files/smpl/SMPL_FEMALE.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d4a1791b6b94880397e1a3a4539b703a228d2150c57de7b288389a8115f4ef0
+size 247530000

common/utils/human_model_files/smpl/SMPL_MALE.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed4d55bb3041fefc6f73b70694d6c8edc1020c0d07340be5cc651cae2c6a6ae3
+size 247101031

common/utils/human_model_files/smpl/SMPL_NEUTRAL.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4924f235e63f7c5d5b690acedf736419c2edb846a2d69fc0956169615fa75688
+size 247186228

common/utils/human_model_files/smpl/smpl_uv.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb2a1aaf8be2091ebc4344daefae0622cc09252b33d4f6c36ea2c6541a01d469
+size 1524004

common/utils/human_model_files/smplx/MANO_SMPLX_vertex_ids.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5abe70b6574de25470475091e8008314a5b90127eb48c3e63bfa0adf8c04dcf
+size 13535

common/utils/human_model_files/smplx/SMPL-X__FLAME_vertex_ids.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e70cdc3659aae699b9732e8dd4af49106310c69b90dc83d9f73e96dbf871e49
+size 40312

common/utils/human_model_files/smplx/SMPLX_FEMALE.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2a3686c9d6d218ff6822fba411c607a3c8125a70af340f384ce68bebecabe0e
+size 108794146

common/utils/human_model_files/smplx/SMPLX_MALE.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab318e3f37d2bfaae26abf4e6fab445c2a610e1d63714794d60379cc263bc2a5
+size 108753445

common/utils/human_model_files/smplx/SMPLX_NEUTRAL.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:376021446ddc86e99acacd795182bbef903e61d33b76b9d8b359c2b0865bd992
+size 108752058

common/utils/human_model_files/smplx/SMPLX_NEUTRAL.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:381c808965deb4f5e845f8c3eddb0cd69930cc72e5774ce4f34c4ce3cf058361
+size 544173380

common/utils/human_model_files/smplx/SMPLX_to_J14.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5df844ddea85b0a400a2e8dbe63d09d19f2b1b7ec0e0e952daeae08f83d82d61
+size 4692193

common/utils/human_models.py

@@ -0,0 +1,176 @@
+import numpy as np
+import torch
+import os.path as osp
+from config import cfg
+from utils.smplx import smplx
+import pickle
+
+class SMPLX(object):
+    def __init__(self):
+        self.layer_arg = {'create_global_orient': False, 'create_body_pose': False, 'create_left_hand_pose': False, 'create_right_hand_pose': False, 'create_jaw_pose': False, 'create_leye_pose': False, 'create_reye_pose': False, 'create_betas': False, 'create_expression': False, 'create_transl': False}
+        self.layer = {'neutral': smplx.create(cfg.human_model_path, 'smplx', gender='NEUTRAL', use_pca=False, use_face_contour=True, **self.layer_arg),
+                        'male': smplx.create(cfg.human_model_path, 'smplx', gender='MALE', use_pca=False, use_face_contour=True, **self.layer_arg),
+                        'female': smplx.create(cfg.human_model_path, 'smplx', gender='FEMALE', use_pca=False, use_face_contour=True, **self.layer_arg)
+                        }
+        self.vertex_num = 10475
+        self.face = self.layer['neutral'].faces
+        self.shape_param_dim = 10
+        self.expr_code_dim = 10
+        with open(osp.join(cfg.human_model_path, 'smplx', 'SMPLX_to_J14.pkl'), 'rb') as f:
+            self.j14_regressor = pickle.load(f, encoding='latin1')
+        with open(osp.join(cfg.human_model_path, 'smplx', 'MANO_SMPLX_vertex_ids.pkl'), 'rb') as f:
+            self.hand_vertex_idx = pickle.load(f, encoding='latin1')
+        self.face_vertex_idx = np.load(osp.join(cfg.human_model_path, 'smplx', 'SMPL-X__FLAME_vertex_ids.npy'))
+        self.J_regressor = self.layer['neutral'].J_regressor.numpy()
+        self.J_regressor_idx = {'pelvis': 0, 'lwrist': 20, 'rwrist': 21, 'neck': 12}
+        self.orig_hand_regressor = self.make_hand_regressor()
+        #self.orig_hand_regressor = {'left': self.layer.J_regressor.numpy()[[20,37,38,39,25,26,27,28,29,30,34,35,36,31,32,33],:], 'right': self.layer.J_regressor.numpy()[[21,52,53,54,40,41,42,43,44,45,49,50,51,46,47,48],:]}
+
+        # original SMPLX joint set
+        self.orig_joint_num = 53 # 22 (body joints) + 30 (hand joints) + 1 (face jaw joint)
+        self.orig_joints_name = \
+        ('Pelvis', 'L_Hip', 'R_Hip', 'Spine_1', 'L_Knee', 'R_Knee', 'Spine_2', 'L_Ankle', 'R_Ankle', 'Spine_3', 'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', # body joints
+        'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', # left hand joints
+        'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', # right hand joints
+        'Jaw' # face jaw joint
+        )
+        self.orig_flip_pairs = \
+        ( (1,2), (4,5), (7,8), (10,11), (13,14), (16,17), (18,19), (20,21), # body joints
+        (22,37), (23,38), (24,39), (25,40), (26,41), (27,42), (28,43), (29,44), (30,45), (31,46), (32,47), (33,48), (34,49), (35,50), (36,51) # hand joints
+        )
+        self.orig_root_joint_idx = self.orig_joints_name.index('Pelvis')
+        self.orig_joint_part = \
+        {'body': range(self.orig_joints_name.index('Pelvis'), self.orig_joints_name.index('R_Wrist')+1),
+        'lhand': range(self.orig_joints_name.index('L_Index_1'), self.orig_joints_name.index('L_Thumb_3')+1),
+        'rhand': range(self.orig_joints_name.index('R_Index_1'), self.orig_joints_name.index('R_Thumb_3')+1),
+        'face': range(self.orig_joints_name.index('Jaw'), self.orig_joints_name.index('Jaw')+1)}
+
+        # changed SMPLX joint set for the supervision
+        self.joint_num = 137 # 25 (body joints) + 40 (hand joints) + 72 (face keypoints)
+        self.joints_name = \
+        ('Pelvis', 'L_Hip', 'R_Hip', 'L_Knee', 'R_Knee', 'L_Ankle', 'R_Ankle', 'Neck', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', 'L_Ear', 'R_Ear', 'L_Eye', 'R_Eye', 'Nose',# body joints
+         'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', 'L_Thumb_4', 'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Index_4', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Middle_4', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Ring_4', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Pinky_4', # left hand joints
+         'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', 'R_Thumb_4', 'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Index_4', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Middle_4', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Ring_4', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Pinky_4', # right hand joints
+         *['Face_' + str(i) for i in range(1,73)] # face keypoints (too many keypoints... omit real names. have same name of keypoints defined in FLAME class)
+         )
+        self.root_joint_idx = self.joints_name.index('Pelvis')
+        self.lwrist_idx = self.joints_name.index('L_Wrist')
+        self.rwrist_idx = self.joints_name.index('R_Wrist')
+        self.neck_idx = self.joints_name.index('Neck')
+        self.flip_pairs = \
+        ( (1,2), (3,4), (5,6), (8,9), (10,11), (12,13), (14,17), (15,18), (16,19), (20,21), (22,23), # body joints
+        (25,45), (26,46), (27,47), (28,48), (29,49), (30,50), (31,51), (32,52), (33,53), (34,54), (35,55), (36,56), (37,57), (38,58), (39,59), (40,60), (41,61), (42,62), (43,63), (44,64), # hand joints
+        (67,68), # face eyeballs
+        (69,78), (70,77), (71,76), (72,75), (73,74), # face eyebrow
+        (83,87), (84,86), # face below nose
+        (88,97), (89,96), (90,95), (91,94), (92,99), (93,98), # face eyes
+        (100,106), (101,105), (102,104), (107,111), (108,110), # face mouth
+        (112,116), (113,115), (117,119), # face lip
+        (120,136), (121,135), (122,134), (123,133), (124,132), (125,131), (126,130), (127,129) # face contours
+        )
+        self.joint_idx = \
+        (0,1,2,4,5,7,8,12,16,17,18,19,20,21,60,61,62,63,64,65,59,58,57,56,55, # body joints
+        37,38,39,66,25,26,27,67,28,29,30,68,34,35,36,69,31,32,33,70, # left hand joints
+        52,53,54,71,40,41,42,72,43,44,45,73,49,50,51,74,46,47,48,75, # right hand joints
+        22,15, # jaw, head
+        57,56, # eyeballs
+        76,77,78,79,80,81,82,83,84,85, # eyebrow
+        86,87,88,89, # nose
+        90,91,92,93,94, # below nose
+        95,96,97,98,99,100,101,102,103,104,105,106, # eyes
+        107, # right mouth
+        108,109,110,111,112, # upper mouth
+        113, # left mouth
+        114,115,116,117,118, # lower mouth
+        119, # right lip
+        120,121,122, # upper lip
+        123, # left lip
+        124,125,126, # lower lip
+        127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143 # face contour
+        )
+        self.joint_part = \
+        {'body': range(self.joints_name.index('Pelvis'), self.joints_name.index('Nose')+1),
+        'lhand': range(self.joints_name.index('L_Thumb_1'), self.joints_name.index('L_Pinky_4')+1),
+        'rhand': range(self.joints_name.index('R_Thumb_1'), self.joints_name.index('R_Pinky_4')+1),
+        'hand': range(self.joints_name.index('L_Thumb_1'), self.joints_name.index('R_Pinky_4')+1),
+        'face': range(self.joints_name.index('Face_1'), self.joints_name.index('Face_72')+1)}
+        
+        # changed SMPLX joint set for PositionNet prediction
+        self.pos_joint_num = 65 # 25 (body joints) + 40 (hand joints)
+        self.pos_joints_name = \
+        ('Pelvis', 'L_Hip', 'R_Hip', 'L_Knee', 'R_Knee', 'L_Ankle', 'R_Ankle', 'Neck', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', 'L_Ear', 'R_Ear', 'L_Eye', 'R_Eye', 'Nose', # body joints
+         'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', 'L_Thumb_4', 'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Index_4', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Middle_4', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Ring_4', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Pinky_4', # left hand joints
+         'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', 'R_Thumb_4', 'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Index_4', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Middle_4', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Ring_4', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Pinky_4', # right hand joints
+         )
+        self.pos_joint_part = \
+        {'body': range(self.pos_joints_name.index('Pelvis'), self.pos_joints_name.index('Nose')+1),
+        'lhand': range(self.pos_joints_name.index('L_Thumb_1'), self.pos_joints_name.index('L_Pinky_4')+1),
+        'rhand': range(self.pos_joints_name.index('R_Thumb_1'), self.pos_joints_name.index('R_Pinky_4')+1),
+        'hand': range(self.pos_joints_name.index('L_Thumb_1'), self.pos_joints_name.index('R_Pinky_4')+1)}
+        self.pos_joint_part['L_MCP'] = [self.pos_joints_name.index('L_Index_1') - len(self.pos_joint_part['body']),
+                                        self.pos_joints_name.index('L_Middle_1') - len(self.pos_joint_part['body']),
+                                        self.pos_joints_name.index('L_Ring_1') - len(self.pos_joint_part['body']),
+                                        self.pos_joints_name.index('L_Pinky_1') - len(self.pos_joint_part['body'])]
+        self.pos_joint_part['R_MCP'] = [self.pos_joints_name.index('R_Index_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand']),
+                                        self.pos_joints_name.index('R_Middle_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand']),
+                                        self.pos_joints_name.index('R_Ring_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand']),
+                                        self.pos_joints_name.index('R_Pinky_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand'])]
+    
+    def make_hand_regressor(self):
+        regressor = self.layer['neutral'].J_regressor.numpy()
+        lhand_regressor = np.concatenate((regressor[[20,37,38,39],:],
+                                            np.eye(self.vertex_num)[5361,None],
+                                                regressor[[25,26,27],:],
+                                                np.eye(self.vertex_num)[4933,None],
+                                                regressor[[28,29,30],:],
+                                                np.eye(self.vertex_num)[5058,None],
+                                                regressor[[34,35,36],:],
+                                                np.eye(self.vertex_num)[5169,None],
+                                                regressor[[31,32,33],:],
+                                                np.eye(self.vertex_num)[5286,None]))
+        rhand_regressor = np.concatenate((regressor[[21,52,53,54],:],
+                                            np.eye(self.vertex_num)[8079,None],
+                                                regressor[[40,41,42],:],
+                                                np.eye(self.vertex_num)[7669,None],
+                                                regressor[[43,44,45],:],
+                                                np.eye(self.vertex_num)[7794,None],
+                                                regressor[[49,50,51],:],
+                                                np.eye(self.vertex_num)[7905,None],
+                                                regressor[[46,47,48],:],
+                                                np.eye(self.vertex_num)[8022,None]))
+        hand_regressor = {'left': lhand_regressor, 'right': rhand_regressor}
+        return hand_regressor
+
+        
+    def reduce_joint_set(self, joint):
+        new_joint = []
+        for name in self.pos_joints_name:
+            idx = self.joints_name.index(name)
+            new_joint.append(joint[:,idx,:])
+        new_joint = torch.stack(new_joint,1)
+        return new_joint
+
+class SMPL(object):
+    def __init__(self):
+        self.layer_arg = {'create_body_pose': False, 'create_betas': False, 'create_global_orient': False, 'create_transl': False}
+        self.layer = {'neutral': smplx.create(cfg.human_model_path, 'smpl', gender='NEUTRAL', **self.layer_arg), 'male': smplx.create(cfg.human_model_path, 'smpl', gender='MALE', **self.layer_arg), 'female': smplx.create(cfg.human_model_path, 'smpl', gender='FEMALE', **self.layer_arg)}
+        self.vertex_num = 6890
+        self.face = self.layer['neutral'].faces
+        self.shape_param_dim = 10
+        self.vposer_code_dim = 32
+
+        # original SMPL joint set
+        self.orig_joint_num = 24
+        self.orig_joints_name = ('Pelvis', 'L_Hip', 'R_Hip', 'Spine_1', 'L_Knee', 'R_Knee', 'Spine_2', 'L_Ankle', 'R_Ankle', 'Spine_3', 'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Hand', 'R_Hand')
+        self.orig_flip_pairs = ( (1,2), (4,5), (7,8), (10,11), (13,14), (16,17), (18,19), (20,21), (22,23) )
+        self.orig_root_joint_idx = self.orig_joints_name.index('Pelvis')
+        self.orig_joint_regressor = self.layer['neutral'].J_regressor.numpy().astype(np.float32)
+        
+        self.joint_num = self.orig_joint_num
+        self.joints_name = self.orig_joints_name
+        self.flip_pairs = self.orig_flip_pairs
+        self.root_joint_idx = self.orig_root_joint_idx
+        self.joint_regressor = self.orig_joint_regressor
+
+smpl_x = SMPLX()
+smpl = SMPL()

common/utils/inference_utils.py

@@ -0,0 +1,153 @@
+from typing import Literal, Union
+
+def process_mmdet_results(mmdet_results: list,
+                          cat_id: int = 0,
+                          multi_person: bool = True) -> list:
+    """Process mmdet results, sort bboxes by area in descending order.
+
+    Args:
+        mmdet_results (list):
+            Result of mmdet.apis.inference_detector
+            when the input is a batch.
+            Shape of the nested lists is
+            (n_frame, n_category, n_human, 5).
+        cat_id (int, optional):
+            Category ID. This function will only select
+            the selected category, and drop the others.
+            Defaults to 0, ID of human category.
+        multi_person (bool, optional):
+            Whether to allow multi-person detection, which is
+            slower than single-person. If false, the function
+            only assure that the first person of each frame
+            has the biggest bbox.
+            Defaults to True.
+
+    Returns:
+        list:
+            A list of detected bounding boxes.
+            Shape of the nested lists is
+            (n_frame, n_human, 5)
+            and each bbox is (x, y, x, y, score).
+    """
+    ret_list = []
+    only_max_arg = not multi_person
+    # for _, frame_results in enumerate(mmdet_results):
+    cat_bboxes = mmdet_results[cat_id]
+    # import pdb; pdb.set_trace()
+    sorted_bbox = qsort_bbox_list(cat_bboxes, only_max_arg)
+
+    if only_max_arg:
+        ret_list.append(sorted_bbox[0:1])
+    else:
+        ret_list.append(sorted_bbox)
+    return ret_list
+
+
+def qsort_bbox_list(bbox_list: list,
+                    only_max: bool = False,
+                    bbox_convention: Literal['xyxy', 'xywh'] = 'xyxy'):
+    """Sort a list of bboxes, by their area in pixel(W*H).
+
+    Args:
+        input_list (list):
+            A list of bboxes. Each item is a list of (x1, y1, x2, y2)
+        only_max (bool, optional):
+            If True, only assure the max element at first place,
+            others may not be well sorted.
+            If False, return a well sorted descending list.
+            Defaults to False.
+        bbox_convention (str, optional):
+            Bbox type, xyxy or xywh. Defaults to 'xyxy'.
+
+    Returns:
+        list:
+            A sorted(maybe not so well) descending list.
+    """
+    # import pdb; pdb.set_trace()
+    if len(bbox_list) <= 1:
+        return bbox_list
+    else:
+        bigger_list = []
+        less_list = []
+        anchor_index = int(len(bbox_list) / 2)
+        anchor_bbox = bbox_list[anchor_index]
+        anchor_area = get_area_of_bbox(anchor_bbox, bbox_convention)
+        for i in range(len(bbox_list)):
+            if i == anchor_index:
+                continue
+            tmp_bbox = bbox_list[i]
+            tmp_area = get_area_of_bbox(tmp_bbox, bbox_convention)
+            if tmp_area >= anchor_area:
+                bigger_list.append(tmp_bbox)
+            else:
+                less_list.append(tmp_bbox)
+        if only_max:
+            return qsort_bbox_list(bigger_list) + \
+                [anchor_bbox, ] + less_list
+        else:
+            return qsort_bbox_list(bigger_list) + \
+                [anchor_bbox, ] + qsort_bbox_list(less_list)
+
+def get_area_of_bbox(
+        bbox: Union[list, tuple],
+        bbox_convention: Literal['xyxy', 'xywh'] = 'xyxy') -> float:
+    """Get the area of a bbox_xyxy.
+
+    Args:
+        (Union[list, tuple]):
+            A list of [x1, y1, x2, y2].
+        bbox_convention (str, optional):
+            Bbox type, xyxy or xywh. Defaults to 'xyxy'.
+
+    Returns:
+        float:
+            Area of the bbox(|y2-y1|*|x2-x1|).
+    """
+    # import pdb;pdb.set_trace()
+    if bbox_convention == 'xyxy':
+        return abs(bbox[2] - bbox[0]) * abs(bbox[3] - bbox[1])
+    elif bbox_convention == 'xywh':
+        return abs(bbox[2] * bbox[3])
+    else:
+        raise TypeError(f'Wrong bbox convention: {bbox_convention}')
+
+def calculate_iou(bbox1, bbox2):
+    # Calculate the Intersection over Union (IoU) between two bounding boxes
+    x1 = max(bbox1[0], bbox2[0])
+    y1 = max(bbox1[1], bbox2[1])
+    x2 = min(bbox1[2], bbox2[2])
+    y2 = min(bbox1[3], bbox2[3])
+    
+    intersection_area = max(0, x2 - x1 + 1) * max(0, y2 - y1 + 1)
+    
+    bbox1_area = (bbox1[2] - bbox1[0] + 1) * (bbox1[3] - bbox1[1] + 1)
+    bbox2_area = (bbox2[2] - bbox2[0] + 1) * (bbox2[3] - bbox2[1] + 1)
+    
+    union_area = bbox1_area + bbox2_area - intersection_area
+    
+    iou = intersection_area / union_area
+    return iou
+
+
+def non_max_suppression(bboxes, iou_threshold):
+    # Sort the bounding boxes by their confidence scores (e.g., the probability of containing an object)
+    bboxes = sorted(bboxes, key=lambda x: x[4], reverse=True)
+    
+    # Initialize a list to store the selected bounding boxes
+    selected_bboxes = []
+    
+    # Perform non-maximum suppression
+    while len(bboxes) > 0:
+        current_bbox = bboxes[0]
+        selected_bboxes.append(current_bbox)
+        bboxes = bboxes[1:]
+        
+        remaining_bboxes = []
+        for bbox in bboxes:
+            iou = calculate_iou(current_bbox, bbox)
+            if iou < iou_threshold:
+                remaining_bboxes.append(bbox)
+                
+        bboxes = remaining_bboxes
+        
+    return selected_bboxes

common/utils/preprocessing.py

@@ -0,0 +1,541 @@
+import numpy as np
+import cv2
+import random
+from config import cfg
+import math
+from utils.human_models import smpl_x, smpl
+from utils.transforms import cam2pixel, transform_joint_to_other_db
+from plyfile import PlyData, PlyElement
+import torch
+
+
+def load_img(path, order='RGB'):
+    img = cv2.imread(path, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
+    if not isinstance(img, np.ndarray):
+        raise IOError("Fail to read %s" % path)
+
+    if order == 'RGB':
+        img = img[:, :, ::-1].copy()
+
+    img = img.astype(np.float32)
+    return img
+
+
+def get_bbox(joint_img, joint_valid, extend_ratio=1.2):
+    x_img, y_img = joint_img[:, 0], joint_img[:, 1]
+    x_img = x_img[joint_valid == 1];
+    y_img = y_img[joint_valid == 1];
+    xmin = min(x_img);
+    ymin = min(y_img);
+    xmax = max(x_img);
+    ymax = max(y_img);
+
+    x_center = (xmin + xmax) / 2.;
+    width = xmax - xmin;
+    xmin = x_center - 0.5 * width * extend_ratio
+    xmax = x_center + 0.5 * width * extend_ratio
+
+    y_center = (ymin + ymax) / 2.;
+    height = ymax - ymin;
+    ymin = y_center - 0.5 * height * extend_ratio
+    ymax = y_center + 0.5 * height * extend_ratio
+
+    bbox = np.array([xmin, ymin, xmax - xmin, ymax - ymin]).astype(np.float32)
+    return bbox
+
+
+def sanitize_bbox(bbox, img_width, img_height):
+    x, y, w, h = bbox
+    x1 = np.max((0, x))
+    y1 = np.max((0, y))
+    x2 = np.min((img_width - 1, x1 + np.max((0, w - 1))))
+    y2 = np.min((img_height - 1, y1 + np.max((0, h - 1))))
+    if w * h > 0 and x2 > x1 and y2 > y1:
+        bbox = np.array([x1, y1, x2 - x1, y2 - y1])
+    else:
+        bbox = None
+
+    return bbox
+
+
+def process_bbox(bbox, img_width, img_height, ratio=1.25):
+    bbox = sanitize_bbox(bbox, img_width, img_height)
+    if bbox is None:
+        return bbox
+
+    # aspect ratio preserving bbox
+    w = bbox[2]
+    h = bbox[3]
+    c_x = bbox[0] + w / 2.
+    c_y = bbox[1] + h / 2.
+    aspect_ratio = cfg.input_img_shape[1] / cfg.input_img_shape[0]
+    if w > aspect_ratio * h:
+        h = w / aspect_ratio
+    elif w < aspect_ratio * h:
+        w = h * aspect_ratio
+    bbox[2] = w * ratio
+    bbox[3] = h * ratio
+    bbox[0] = c_x - bbox[2] / 2.
+    bbox[1] = c_y - bbox[3] / 2.
+
+    bbox = bbox.astype(np.float32)
+    return bbox
+
+
+def get_aug_config():
+    scale_factor = 0.25
+    rot_factor = 30
+    color_factor = 0.2
+
+    scale = np.clip(np.random.randn(), -1.0, 1.0) * scale_factor + 1.0
+    rot = np.clip(np.random.randn(), -2.0,
+                  2.0) * rot_factor if random.random() <= 0.6 else 0
+    c_up = 1.0 + color_factor
+    c_low = 1.0 - color_factor
+    color_scale = np.array([random.uniform(c_low, c_up), random.uniform(c_low, c_up), random.uniform(c_low, c_up)])
+    do_flip = random.random() <= 0.5
+
+    return scale, rot, color_scale, do_flip
+
+
+def augmentation(img, bbox, data_split):
+    if getattr(cfg, 'no_aug', False):
+        scale, rot, color_scale, do_flip = 1.0, 0.0, np.array([1, 1, 1]), False
+    elif data_split == 'train':
+        scale, rot, color_scale, do_flip = get_aug_config()
+    else:
+        scale, rot, color_scale, do_flip = 1.0, 0.0, np.array([1, 1, 1]), False
+
+    img, trans, inv_trans = generate_patch_image(img, bbox, scale, rot, do_flip, cfg.input_img_shape)
+    img = np.clip(img * color_scale[None, None, :], 0, 255)
+    return img, trans, inv_trans, rot, do_flip
+
+
+def generate_patch_image(cvimg, bbox, scale, rot, do_flip, out_shape):
+    img = cvimg.copy()
+    img_height, img_width, img_channels = img.shape
+
+    bb_c_x = float(bbox[0] + 0.5 * bbox[2])
+    bb_c_y = float(bbox[1] + 0.5 * bbox[3])
+    bb_width = float(bbox[2])
+    bb_height = float(bbox[3])
+
+    if do_flip:
+        img = img[:, ::-1, :]
+        bb_c_x = img_width - bb_c_x - 1
+
+    trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot)
+    img_patch = cv2.warpAffine(img, trans, (int(out_shape[1]), int(out_shape[0])), flags=cv2.INTER_LINEAR)
+    img_patch = img_patch.astype(np.float32)
+    inv_trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot,
+                                        inv=True)
+
+    return img_patch, trans, inv_trans
+
+
+def rotate_2d(pt_2d, rot_rad):
+    x = pt_2d[0]
+    y = pt_2d[1]
+    sn, cs = np.sin(rot_rad), np.cos(rot_rad)
+    xx = x * cs - y * sn
+    yy = x * sn + y * cs
+    return np.array([xx, yy], dtype=np.float32)
+
+
+def gen_trans_from_patch_cv(c_x, c_y, src_width, src_height, dst_width, dst_height, scale, rot, inv=False):
+    # augment size with scale
+    src_w = src_width * scale
+    src_h = src_height * scale
+    src_center = np.array([c_x, c_y], dtype=np.float32)
+
+    # augment rotation
+    rot_rad = np.pi * rot / 180
+    src_downdir = rotate_2d(np.array([0, src_h * 0.5], dtype=np.float32), rot_rad)
+    src_rightdir = rotate_2d(np.array([src_w * 0.5, 0], dtype=np.float32), rot_rad)
+
+    dst_w = dst_width
+    dst_h = dst_height
+    dst_center = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32)
+    dst_downdir = np.array([0, dst_h * 0.5], dtype=np.float32)
+    dst_rightdir = np.array([dst_w * 0.5, 0], dtype=np.float32)
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = src_center
+    src[1, :] = src_center + src_downdir
+    src[2, :] = src_center + src_rightdir
+
+    dst = np.zeros((3, 2), dtype=np.float32)
+    dst[0, :] = dst_center
+    dst[1, :] = dst_center + dst_downdir
+    dst[2, :] = dst_center + dst_rightdir
+
+    if inv:
+        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    trans = trans.astype(np.float32)
+    return trans
+
+
+def process_db_coord(joint_img, joint_cam, joint_valid, do_flip, img_shape, flip_pairs, img2bb_trans, rot,
+                     src_joints_name, target_joints_name):
+    joint_img_original = joint_img.copy()
+    joint_img, joint_cam, joint_valid = joint_img.copy(), joint_cam.copy(), joint_valid.copy()
+
+    # flip augmentation
+    if do_flip:
+        joint_cam[:, 0] = -joint_cam[:, 0]
+        joint_img[:, 0] = img_shape[1] - 1 - joint_img[:, 0]
+        for pair in flip_pairs:
+            joint_img[pair[0], :], joint_img[pair[1], :] = joint_img[pair[1], :].copy(), joint_img[pair[0], :].copy()
+            joint_cam[pair[0], :], joint_cam[pair[1], :] = joint_cam[pair[1], :].copy(), joint_cam[pair[0], :].copy()
+            joint_valid[pair[0], :], joint_valid[pair[1], :] = joint_valid[pair[1], :].copy(), joint_valid[pair[0],
+                                                                                               :].copy()
+
+    # 3D data rotation augmentation
+    rot_aug_mat = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0],
+                            [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],
+                            [0, 0, 1]], dtype=np.float32)
+    joint_cam = np.dot(rot_aug_mat, joint_cam.transpose(1, 0)).transpose(1, 0)
+
+    # affine transformation
+    joint_img_xy1 = np.concatenate((joint_img[:, :2], np.ones_like(joint_img[:, :1])), 1)
+    joint_img[:, :2] = np.dot(img2bb_trans, joint_img_xy1.transpose(1, 0)).transpose(1, 0)
+    joint_img[:, 0] = joint_img[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
+    joint_img[:, 1] = joint_img[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]
+
+    # check truncation
+    joint_trunc = joint_valid * ((joint_img_original[:, 0] > 0) * (joint_img[:, 0] >= 0) * (joint_img[:, 0] < cfg.output_hm_shape[2]) * \
+                                 (joint_img_original[:, 1] > 0) *(joint_img[:, 1] >= 0) * (joint_img[:, 1] < cfg.output_hm_shape[1]) * \
+                                 (joint_img_original[:, 2] > 0) *(joint_img[:, 2] >= 0) * (joint_img[:, 2] < cfg.output_hm_shape[0])).reshape(-1,
+                                                                                                              1).astype(
+        np.float32)
+
+    # transform joints to target db joints
+    joint_img = transform_joint_to_other_db(joint_img, src_joints_name, target_joints_name)
+    joint_cam_wo_ra = transform_joint_to_other_db(joint_cam, src_joints_name, target_joints_name)
+    joint_valid = transform_joint_to_other_db(joint_valid, src_joints_name, target_joints_name)
+    joint_trunc = transform_joint_to_other_db(joint_trunc, src_joints_name, target_joints_name)
+
+    # root-alignment, for joint_cam input wo ra
+    joint_cam_ra = joint_cam_wo_ra.copy()
+    joint_cam_ra = joint_cam_ra - joint_cam_ra[smpl_x.root_joint_idx, None, :]  # root-relative
+    joint_cam_ra[smpl_x.joint_part['lhand'], :] = joint_cam_ra[smpl_x.joint_part['lhand'], :] - joint_cam_ra[
+                                                                                            smpl_x.lwrist_idx, None,
+                                                                                            :]  # left hand root-relative
+    joint_cam_ra[smpl_x.joint_part['rhand'], :] = joint_cam_ra[smpl_x.joint_part['rhand'], :] - joint_cam_ra[
+                                                                                            smpl_x.rwrist_idx, None,
+                                                                                            :]  # right hand root-relative
+    joint_cam_ra[smpl_x.joint_part['face'], :] = joint_cam_ra[smpl_x.joint_part['face'], :] - joint_cam_ra[smpl_x.neck_idx,
+                                                                                        None,
+                                                                                        :]  # face root-relative
+
+    return joint_img, joint_cam_wo_ra, joint_cam_ra, joint_valid, joint_trunc
+
+
+def process_human_model_output(human_model_param, cam_param, do_flip, img_shape, img2bb_trans, rot, human_model_type, joint_img=None):
+    if human_model_type == 'smplx':
+        human_model = smpl_x
+        rotation_valid = np.ones((smpl_x.orig_joint_num), dtype=np.float32)
+        coord_valid = np.ones((smpl_x.joint_num), dtype=np.float32)
+
+        root_pose, body_pose, shape, trans = human_model_param['root_pose'], human_model_param['body_pose'], \
+                                             human_model_param['shape'], human_model_param['trans']
+        if 'lhand_pose' in human_model_param and human_model_param['lhand_valid']:
+            lhand_pose = human_model_param['lhand_pose']
+        else:
+            lhand_pose = np.zeros((3 * len(smpl_x.orig_joint_part['lhand'])), dtype=np.float32)
+            rotation_valid[smpl_x.orig_joint_part['lhand']] = 0
+            coord_valid[smpl_x.joint_part['lhand']] = 0
+        if 'rhand_pose' in human_model_param and human_model_param['rhand_valid']:
+            rhand_pose = human_model_param['rhand_pose']
+        else:
+            rhand_pose = np.zeros((3 * len(smpl_x.orig_joint_part['rhand'])), dtype=np.float32)
+            rotation_valid[smpl_x.orig_joint_part['rhand']] = 0
+            coord_valid[smpl_x.joint_part['rhand']] = 0
+        if 'jaw_pose' in human_model_param and 'expr' in human_model_param and human_model_param['face_valid']:
+            jaw_pose = human_model_param['jaw_pose']
+            expr = human_model_param['expr']
+            expr_valid = True
+        else:
+            jaw_pose = np.zeros((3), dtype=np.float32)
+            expr = np.zeros((smpl_x.expr_code_dim), dtype=np.float32)
+            rotation_valid[smpl_x.orig_joint_part['face']] = 0
+            coord_valid[smpl_x.joint_part['face']] = 0
+            expr_valid = False
+        if 'gender' in human_model_param:
+            gender = human_model_param['gender']
+        else:
+            gender = 'neutral'
+        root_pose = torch.FloatTensor(root_pose).view(1, 3)  # (1,3)
+        body_pose = torch.FloatTensor(body_pose).view(-1, 3)  # (21,3)
+        lhand_pose = torch.FloatTensor(lhand_pose).view(-1, 3)  # (15,3)
+        rhand_pose = torch.FloatTensor(rhand_pose).view(-1, 3)  # (15,3)
+        jaw_pose = torch.FloatTensor(jaw_pose).view(-1, 3)  # (1,3)
+        shape = torch.FloatTensor(shape).view(1, -1)  # SMPLX shape parameter
+        expr = torch.FloatTensor(expr).view(1, -1)  # SMPLX expression parameter
+        trans = torch.FloatTensor(trans).view(1, -1)  # translation vector
+
+        # apply camera extrinsic (rotation)
+        # merge root pose and camera rotation
+        if 'R' in cam_param:
+            R = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3)
+            root_pose = root_pose.numpy()
+            root_pose, _ = cv2.Rodrigues(root_pose)
+            root_pose, _ = cv2.Rodrigues(np.dot(R, root_pose))
+            root_pose = torch.from_numpy(root_pose).view(1, 3)
+
+        # get mesh and joint coordinates
+        zero_pose = torch.zeros((1, 3)).float()  # eye poses
+        with torch.no_grad():
+            output = smpl_x.layer[gender](betas=shape, body_pose=body_pose.view(1, -1), global_orient=root_pose,
+                                          transl=trans, left_hand_pose=lhand_pose.view(1, -1),
+                                          right_hand_pose=rhand_pose.view(1, -1), jaw_pose=jaw_pose.view(1, -1),
+                                          leye_pose=zero_pose, reye_pose=zero_pose, expression=expr)
+        mesh_cam = output.vertices[0].numpy()
+        joint_cam = output.joints[0].numpy()[smpl_x.joint_idx, :]
+
+        # apply camera exrinsic (translation)
+        # compenstate rotation (translation from origin to root joint was not cancled)
+        if 'R' in cam_param and 't' in cam_param:
+            R, t = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3), np.array(cam_param['t'],
+                                                                                      dtype=np.float32).reshape(1, 3)
+            root_cam = joint_cam[smpl_x.root_joint_idx, None, :]
+            joint_cam = joint_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
+            mesh_cam = mesh_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
+
+        # concat root, body, two hands, and jaw pose
+        pose = torch.cat((root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose))
+
+        # joint coordinates
+        if 'focal' not in cam_param or 'princpt' not in cam_param:
+            assert joint_img is not None 
+        else:   
+            joint_img = cam2pixel(joint_cam, cam_param['focal'], cam_param['princpt'])
+
+        joint_img_original = joint_img.copy()
+
+        joint_cam = joint_cam - joint_cam[smpl_x.root_joint_idx, None, :]  # root-relative
+        joint_cam[smpl_x.joint_part['lhand'], :] = joint_cam[smpl_x.joint_part['lhand'], :] - joint_cam[
+                                                                                              smpl_x.lwrist_idx, None,
+                                                                                              :]  # left hand root-relative
+        joint_cam[smpl_x.joint_part['rhand'], :] = joint_cam[smpl_x.joint_part['rhand'], :] - joint_cam[
+                                                                                              smpl_x.rwrist_idx, None,
+                                                                                              :]  # right hand root-relative
+        joint_cam[smpl_x.joint_part['face'], :] = joint_cam[smpl_x.joint_part['face'], :] - joint_cam[smpl_x.neck_idx,
+                                                                                            None,
+                                                                                            :]  # face root-relative
+        joint_img[smpl_x.joint_part['body'], 2] = (joint_cam[smpl_x.joint_part['body'], 2].copy() / (
+                    cfg.body_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0]  # body depth discretize
+        joint_img[smpl_x.joint_part['lhand'], 2] = (joint_cam[smpl_x.joint_part['lhand'], 2].copy() / (
+                    cfg.hand_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0]  # left hand depth discretize
+        joint_img[smpl_x.joint_part['rhand'], 2] = (joint_cam[smpl_x.joint_part['rhand'], 2].copy() / (
+                    cfg.hand_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0]  # right hand depth discretize
+        joint_img[smpl_x.joint_part['face'], 2] = (joint_cam[smpl_x.joint_part['face'], 2].copy() / (
+                    cfg.face_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0]  # face depth discretize
+
+    elif human_model_type == 'smpl':
+        human_model = smpl
+        pose, shape, trans = human_model_param['pose'], human_model_param['shape'], human_model_param['trans']
+        if 'gender' in human_model_param:
+            gender = human_model_param['gender']
+        else:
+            gender = 'neutral'
+        pose = torch.FloatTensor(pose).view(-1, 3)
+        shape = torch.FloatTensor(shape).view(1, -1);
+        trans = torch.FloatTensor(trans).view(1, -1)  # translation vector
+
+        # apply camera extrinsic (rotation)
+        # merge root pose and camera rotation
+        if 'R' in cam_param:
+            R = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3)
+            root_pose = pose[smpl.orig_root_joint_idx, :].numpy()
+            root_pose, _ = cv2.Rodrigues(root_pose)
+            root_pose, _ = cv2.Rodrigues(np.dot(R, root_pose))
+            pose[smpl.orig_root_joint_idx] = torch.from_numpy(root_pose).view(3)
+
+        # get mesh and joint coordinates
+        root_pose = pose[smpl.orig_root_joint_idx].view(1, 3)
+        body_pose = torch.cat((pose[:smpl.orig_root_joint_idx, :], pose[smpl.orig_root_joint_idx + 1:, :])).view(1, -1)
+        with torch.no_grad():
+            output = smpl.layer[gender](betas=shape, body_pose=body_pose, global_orient=root_pose, transl=trans)
+        mesh_cam = output.vertices[0].numpy()
+        joint_cam = np.dot(smpl.joint_regressor, mesh_cam)
+
+        # apply camera exrinsic (translation)
+        # compenstate rotation (translation from origin to root joint was not cancled)
+        if 'R' in cam_param and 't' in cam_param:
+            R, t = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3), np.array(cam_param['t'],
+                                                                                      dtype=np.float32).reshape(1, 3)
+            root_cam = joint_cam[smpl.root_joint_idx, None, :]
+            joint_cam = joint_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
+            mesh_cam = mesh_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
+
+        # joint coordinates
+        if 'focal' not in cam_param or 'princpt' not in cam_param:
+            assert joint_img is not None 
+        else:   
+            joint_img = cam2pixel(joint_cam, cam_param['focal'], cam_param['princpt'])
+        
+        joint_img_original = joint_img.copy()
+        joint_cam = joint_cam - joint_cam[smpl.root_joint_idx, None, :]  # body root-relative
+        joint_img[:, 2] = (joint_cam[:, 2].copy() / (cfg.body_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[
+            0]  # body depth discretize
+
+    elif human_model_type == 'mano':
+        human_model = mano
+        pose, shape, trans = human_model_param['pose'], human_model_param['shape'], human_model_param['trans']
+        hand_type = human_model_param['hand_type']
+        pose = torch.FloatTensor(pose).view(-1, 3)
+        shape = torch.FloatTensor(shape).view(1, -1);
+        trans = torch.FloatTensor(trans).view(1, -1)  # translation vector
+
+        # apply camera extrinsic (rotation)
+        # merge root pose and camera rotation
+        if 'R' in cam_param:
+            R = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3)
+            root_pose = pose[mano.orig_root_joint_idx, :].numpy()
+            root_pose, _ = cv2.Rodrigues(root_pose)
+            root_pose, _ = cv2.Rodrigues(np.dot(R, root_pose))
+            pose[mano.orig_root_joint_idx] = torch.from_numpy(root_pose).view(3)
+
+        # get mesh and joint coordinates
+        root_pose = pose[mano.orig_root_joint_idx].view(1, 3)
+        hand_pose = torch.cat((pose[:mano.orig_root_joint_idx, :], pose[mano.orig_root_joint_idx + 1:, :])).view(1, -1)
+        with torch.no_grad():
+            output = mano.layer[hand_type](betas=shape, hand_pose=hand_pose, global_orient=root_pose, transl=trans)
+        mesh_cam = output.vertices[0].numpy()
+        joint_cam = np.dot(mano.joint_regressor, mesh_cam)
+
+        # apply camera exrinsic (translation)
+        # compenstate rotation (translation from origin to root joint was not cancled)
+        if 'R' in cam_param and 't' in cam_param:
+            R, t = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3), np.array(cam_param['t'],
+                                                                                      dtype=np.float32).reshape(1, 3)
+            root_cam = joint_cam[mano.root_joint_idx, None, :]
+            joint_cam = joint_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
+            mesh_cam = mesh_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
+
+        # joint coordinates
+        if 'focal' not in cam_param or 'princpt' not in cam_param:
+            assert joint_img is not None 
+        else:   
+            joint_img = cam2pixel(joint_cam, cam_param['focal'], cam_param['princpt'])
+        joint_cam = joint_cam - joint_cam[mano.root_joint_idx, None, :]  # hand root-relative
+        joint_img[:, 2] = (joint_cam[:, 2].copy() / (cfg.hand_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[
+            0]  # hand depth discretize
+
+    mesh_cam_orig = mesh_cam.copy()  # back-up the original one
+
+    ## so far, data augmentations are not applied yet
+    ## now, apply data augmentations
+
+    # image projection
+    if do_flip:
+        joint_cam[:, 0] = -joint_cam[:, 0]
+        joint_img[:, 0] = img_shape[1] - 1 - joint_img[:, 0]
+        for pair in human_model.flip_pairs:
+            joint_cam[pair[0], :], joint_cam[pair[1], :] = joint_cam[pair[1], :].copy(), joint_cam[pair[0], :].copy()
+            joint_img[pair[0], :], joint_img[pair[1], :] = joint_img[pair[1], :].copy(), joint_img[pair[0], :].copy()
+            if human_model_type == 'smplx':
+                coord_valid[pair[0]], coord_valid[pair[1]] = coord_valid[pair[1]].copy(), coord_valid[pair[0]].copy()
+
+    # x,y affine transform, root-relative depth
+    joint_img_xy1 = np.concatenate((joint_img[:, :2], np.ones_like(joint_img[:, 0:1])), 1)
+    joint_img[:, :2] = np.dot(img2bb_trans, joint_img_xy1.transpose(1, 0)).transpose(1, 0)[:, :2]
+    joint_img[:, 0] = joint_img[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
+    joint_img[:, 1] = joint_img[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]
+
+    # check truncation
+    # TODO
+    joint_trunc = ((joint_img_original[:, 0] > 0) * (joint_img[:, 0] >= 0) * (joint_img[:, 0] < cfg.output_hm_shape[2]) * \
+                   (joint_img_original[:, 1] > 0) * (joint_img[:, 1] >= 0) * (joint_img[:, 1] < cfg.output_hm_shape[1]) * \
+                   (joint_img_original[:, 2] > 0) * (joint_img[:, 2] >= 0) * (joint_img[:, 2] < cfg.output_hm_shape[0])).reshape(-1, 1).astype(
+        np.float32)
+
+    # 3D data rotation augmentation
+    rot_aug_mat = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0],
+                            [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],
+                            [0, 0, 1]], dtype=np.float32)
+    # coordinate
+    joint_cam = np.dot(rot_aug_mat, joint_cam.transpose(1, 0)).transpose(1, 0)
+    # parameters
+    # flip pose parameter (axis-angle)
+    if do_flip:
+        for pair in human_model.orig_flip_pairs:
+            pose[pair[0], :], pose[pair[1], :] = pose[pair[1], :].clone(), pose[pair[0], :].clone()
+            if human_model_type == 'smplx':
+                rotation_valid[pair[0]], rotation_valid[pair[1]] = rotation_valid[pair[1]].copy(), rotation_valid[
+                    pair[0]].copy()
+        pose[:, 1:3] *= -1  # multiply -1 to y and z axis of axis-angle
+
+    # rotate root pose
+    pose = pose.numpy()
+    root_pose = pose[human_model.orig_root_joint_idx, :]
+    root_pose, _ = cv2.Rodrigues(root_pose)
+    root_pose, _ = cv2.Rodrigues(np.dot(rot_aug_mat, root_pose))
+    pose[human_model.orig_root_joint_idx] = root_pose.reshape(3)
+
+    # change to mean shape if beta is too far from it
+    shape[(shape.abs() > 3).any(dim=1)] = 0.
+    shape = shape.numpy().reshape(-1)
+
+    # return results
+    if human_model_type == 'smplx':
+        pose = pose.reshape(-1)
+        expr = expr.numpy().reshape(-1)
+
+        return joint_img, joint_cam, joint_trunc, pose, shape, expr, rotation_valid, coord_valid, expr_valid, mesh_cam_orig
+    elif human_model_type == 'smpl':
+        pose = pose.reshape(-1)
+        return joint_img, joint_cam, joint_trunc, pose, shape, mesh_cam_orig
+    elif human_model_type == 'mano':
+        pose = pose.reshape(-1)
+        return joint_img, joint_cam, joint_trunc, pose, shape, mesh_cam_orig
+
+
+def get_fitting_error_3D(db_joint, db_joint_from_fit, joint_valid):
+    # mask coordinate
+    db_joint = db_joint[np.tile(joint_valid, (1, 3)) == 1].reshape(-1, 3)
+    db_joint_from_fit = db_joint_from_fit[np.tile(joint_valid, (1, 3)) == 1].reshape(-1, 3)
+
+    db_joint_from_fit = db_joint_from_fit - np.mean(db_joint_from_fit, 0)[None, :] + np.mean(db_joint, 0)[None,
+                                                                                     :]  # translation alignment
+    error = np.sqrt(np.sum((db_joint - db_joint_from_fit) ** 2, 1)).mean()
+    return error
+
+
+def load_obj(file_name):
+    v = []
+    obj_file = open(file_name)
+    for line in obj_file:
+        words = line.split(' ')
+        if words[0] == 'v':
+            x, y, z = float(words[1]), float(words[2]), float(words[3])
+            v.append(np.array([x, y, z]))
+    return np.stack(v)
+
+
+def load_ply(file_name):
+    plydata = PlyData.read(file_name)
+    x = plydata['vertex']['x']
+    y = plydata['vertex']['y']
+    z = plydata['vertex']['z']
+    v = np.stack((x, y, z), 1)
+    return v
+
+def resize_bbox(bbox, scale=1.2):
+    if isinstance(bbox, list):
+        x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
+    else:
+        x1, y1, x2, y2 = bbox
+    x_center = (x1+x2)/2.0
+    y_center = (y1+y2)/2.0
+    x_size, y_size = x2-x1, y2-y1
+    x1_resize = x_center-x_size/2.0*scale
+    x2_resize = x_center+x_size/2.0*scale
+    y1_resize = y_center - y_size / 2.0 * scale
+    y2_resize = y_center + y_size / 2.0 * scale
+    bbox[0], bbox[1], bbox[2], bbox[3] = x1_resize, y1_resize, x2_resize, y2_resize
+    return bbox

common/utils/smplx/LICENSE

@@ -0,0 +1,58 @@
+License
+
+Software Copyright License for non-commercial scientific research purposes
+Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License
+
+Ownership / Licensees
+The Software and the associated materials has been developed at the
+
+Max Planck Institute for Intelligent Systems (hereinafter "MPI").
+
+Any copyright or patent right is owned by and proprietary material of the
+
+Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”)
+
+hereinafter the “Licensor”.
+
+License Grant
+Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right:
+
+To install the Model & Software on computers owned, leased or otherwise controlled by you and/or your organization;
+To use the Model & Software for the sole purpose of performing non-commercial scientific research, non-commercial education, or non-commercial artistic projects;
+Any other use, in particular any use for commercial purposes, is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artifacts for commercial purposes. The Model & Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission.
+
+The Model & Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Model & Software to train methods/algorithms/neural networks/etc. for commercial use of any kind. By downloading the Model & Software, you agree not to reverse engineer it.
+
+No Distribution
+The Model & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only.
+
+Disclaimer of Representations and Warranties
+You expressly acknowledge and agree that the Model & Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Model & Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE MODEL & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Model & Software, (ii) that the use of the Model & Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Model & Software will not cause any damage of any kind to you or a third party.
+
+Limitation of Liability
+Because this Model & Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage.
+Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded.
+Patent claims generated through the usage of the Model & Software cannot be directed towards the copyright holders.
+The Model & Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Model & Software and is not responsible for any problems such modifications cause.
+
+No Maintenance Services
+You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Model & Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Model & Software at any time.
+
+Defects of the Model & Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication.
+
+Publications using the Model & Software
+You acknowledge that the Model & Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Model & Software.
+
+Citation:
+
+
+@inproceedings{SMPL-X:2019,
+  title = {Expressive Body Capture: 3D Hands, Face, and Body from a Single Image},
+  author = {Pavlakos, Georgios and Choutas, Vasileios and Ghorbani, Nima and Bolkart, Timo and Osman, Ahmed A. A. and Tzionas, Dimitrios and Black, Michael J.},
+  booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)},
+  year = {2019}
+}
+Commercial licensing opportunities
+For commercial uses of the Software, please send email to [email protected]
+
+This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention.

common/utils/smplx/README.md

@@ -0,0 +1,186 @@
+## SMPL-X:  A new joint 3D model of the human body, face and hands together
+
+[[Paper Page](https://smpl-x.is.tue.mpg.de)] [[Paper](https://ps.is.tuebingen.mpg.de/uploads_file/attachment/attachment/497/SMPL-X.pdf)]
+[[Supp. Mat.](https://ps.is.tuebingen.mpg.de/uploads_file/attachment/attachment/498/SMPL-X-supp.pdf)]
+
+![SMPL-X Examples](./images/teaser_fig.png)
+
+## Table of Contents
+  * [License](#license)
+  * [Description](#description)
+  * [Installation](#installation)
+  * [Downloading the model](#downloading-the-model)
+  * [Loading SMPL-X, SMPL+H and SMPL](#loading-smpl-x-smplh-and-smpl) 
+    * [SMPL and SMPL+H setup](#smpl-and-smplh-setup)
+    * [Model loading](https://github.com/vchoutas/smplx#model-loading)
+  * [MANO and FLAME correspondences](#mano-and-flame-correspondences) 
+  * [Example](#example)
+  * [Citation](#citation)
+  * [Acknowledgments](#acknowledgments)
+  * [Contact](#contact)
+
+## License
+
+Software Copyright License for **non-commercial scientific research purposes**.
+Please read carefully the [terms and conditions](https://github.com/vchoutas/smplx/blob/master/LICENSE) and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this [License](./LICENSE).
+
+## Disclaimer
+
+The original images used for the figures 1 and 2 of the paper can be found in this link. 
+The images in the paper are used under license from gettyimages.com.
+We have acquired the right to use them in the publication, but redistribution is not allowed.
+Please follow the instructions on the given link to acquire right of usage.
+Our results are obtained on the 483 × 724 pixels resolution of the original images.
+
+## Description
+
+*SMPL-X* (SMPL eXpressive) is a unified body model with shape parameters trained jointly for the
+face, hands and body. *SMPL-X* uses standard vertex based linear blend skinning with learned corrective blend
+shapes, has N = 10, 475 vertices and K = 54 joints,
+which include joints for the neck, jaw, eyeballs and fingers. 
+SMPL-X is defined by a function M(θ, β, ψ), where θ is the pose parameters, β the shape parameters and
+ψ the facial expression parameters.
+
+
+## Installation
+
+To install the model please follow the next steps in the specified order:
+1. To install from PyPi simply run: 
+  ```Shell
+  pip install smplx[all]
+  ```
+2. Clone this repository and install it using the *setup.py* script: 
+```Shell
+git clone https://github.com/vchoutas/smplx
+python setup.py install
+```
+
+## Downloading the model
+
+To download the *SMPL-X* model go to [this project website](https://smpl-x.is.tue.mpg.de) and register to get access to the downloads section. 
+
+To download the *SMPL+H* model go to [this project website](http://mano.is.tue.mpg.de) and register to get access to the downloads section. 
+
+To download the *SMPL* model go to [this](http://smpl.is.tue.mpg.de) (male and female models) and [this](http://smplify.is.tue.mpg.de) (gender neutral model) project website and register to get access to the downloads section. 
+
+## Loading SMPL-X, SMPL+H and SMPL
+
+### SMPL and SMPL+H setup
+
+The loader gives the option to use any of the SMPL-X, SMPL+H, SMPL, and MANO models. Depending on the model you want to use, please follow the respective download instructions. To switch between MANO, SMPL, SMPL+H and SMPL-X just change the *model_path* or *model_type* parameters. For more details please check the docs of the model classes.
+Before using SMPL and SMPL+H you should follow the instructions in [tools/README.md](./tools/README.md) to remove the
+Chumpy objects from both model pkls, as well as merge the MANO parameters with SMPL+H.
+
+### Model loading 
+
+You can either use the [create](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L54)
+function from [body_models](./smplx/body_models.py) or directly call the constructor for the 
+[SMPL](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L106), 
+[SMPL+H](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L395) and 
+[SMPL-X](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L628) model. The path to the model can either be the path to the file with the parameters or a directory with the following structure:
+```bash
+models
+├── smpl
+│   ├── SMPL_FEMALE.pkl
+│   ��── SMPL_MALE.pkl
+│   └── SMPL_NEUTRAL.pkl
+├── smplh
+│   ├── SMPLH_FEMALE.pkl
+│   └── SMPLH_MALE.pkl
+├── mano
+|   ├── MANO_RIGHT.pkl
+|   └── MANO_LEFT.pkl
+└── smplx
+    ├── SMPLX_FEMALE.npz
+    ├── SMPLX_FEMALE.pkl
+    ├── SMPLX_MALE.npz
+    ├── SMPLX_MALE.pkl
+    ├── SMPLX_NEUTRAL.npz
+    └── SMPLX_NEUTRAL.pkl
+```
+
+
+## MANO and FLAME correspondences
+
+The vertex correspondences between SMPL-X and MANO, FLAME can be downloaded
+from [the project website](https://smpl-x.is.tue.mpg.de). If you have extracted
+the correspondence data in the folder *correspondences*, then use the following
+scripts to visualize them:
+
+1. To view MANO correspondences run the following command:
+
+```
+python examples/vis_mano_vertices.py --model-folder $SMPLX_FOLDER --corr-fname correspondences/MANO_SMPLX_vertex_ids.pkl
+```
+
+2. To view FLAME correspondences run the following command:
+
+```
+python examples/vis_flame_vertices.py --model-folder $SMPLX_FOLDER --corr-fname correspondences/SMPL-X__FLAME_vertex_ids.npy
+```
+
+## Example
+
+After installing the *smplx* package and downloading the model parameters you should be able to run the *demo.py*
+script to visualize the results. For this step you have to install the [pyrender](https://pyrender.readthedocs.io/en/latest/index.html) and [trimesh](https://trimsh.org/) packages.
+
+`python examples/demo.py --model-folder $SMPLX_FOLDER --plot-joints=True --gender="neutral"`
+
+![SMPL-X Examples](./images/example.png)
+
+## Citation
+
+Depending on which model is loaded for your project, i.e. SMPL-X or SMPL+H or SMPL, please cite the most relevant work below, listed in the same order:
+
+```
+@inproceedings{SMPL-X:2019,
+    title = {Expressive Body Capture: 3D Hands, Face, and Body from a Single Image},
+    author = {Pavlakos, Georgios and Choutas, Vasileios and Ghorbani, Nima and Bolkart, Timo and Osman, Ahmed A. A. and Tzionas, Dimitrios and Black, Michael J.},
+    booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)},
+    year = {2019}
+}
+```
+
+```
+@article{MANO:SIGGRAPHASIA:2017,
+    title = {Embodied Hands: Modeling and Capturing Hands and Bodies Together},
+    author = {Romero, Javier and Tzionas, Dimitrios and Black, Michael J.},
+    journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)},
+    volume = {36},
+    number = {6},
+    series = {245:1--245:17},
+    month = nov,
+    year = {2017},
+    month_numeric = {11}
+  }
+```
+
+```
+@article{SMPL:2015,
+    author = {Loper, Matthew and Mahmood, Naureen and Romero, Javier and Pons-Moll, Gerard and Black, Michael J.},
+    title = {{SMPL}: A Skinned Multi-Person Linear Model},
+    journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)},
+    month = oct,
+    number = {6},
+    pages = {248:1--248:16},
+    publisher = {ACM},
+    volume = {34},
+    year = {2015}
+}
+```
+
+This repository was originally developed for SMPL-X / SMPLify-X (CVPR 2019), you might be interested in having a look: [https://smpl-x.is.tue.mpg.de](https://smpl-x.is.tue.mpg.de).
+
+## Acknowledgments
+
+### Facial Contour
+
+Special thanks to [Soubhik Sanyal](https://github.com/soubhiksanyal) for sharing the Tensorflow code used for the facial
+landmarks.
+
+## Contact
+The code of this repository was implemented by [Vassilis Choutas]([email protected]).
+
+For questions, please contact [[email protected]]([email protected]).
+
+For commercial licensing (and all related questions for business applications), please contact [[email protected]]([email protected]).

common/utils/smplx/examples/demo.py

@@ -0,0 +1,180 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os.path as osp
+import argparse
+
+import numpy as np
+import torch
+
+import smplx
+
+
+def main(model_folder,
+         model_type='smplx',
+         ext='npz',
+         gender='neutral',
+         plot_joints=False,
+         num_betas=10,
+         sample_shape=True,
+         sample_expression=True,
+         num_expression_coeffs=10,
+         plotting_module='pyrender',
+         use_face_contour=False):
+
+    model = smplx.create(model_folder, model_type=model_type,
+                         gender=gender, use_face_contour=use_face_contour,
+                         num_betas=num_betas,
+                         num_expression_coeffs=num_expression_coeffs,
+                         ext=ext)
+    print(model)
+
+    betas, expression = None, None
+    if sample_shape:
+        betas = torch.randn([1, model.num_betas], dtype=torch.float32)
+    if sample_expression:
+        expression = torch.randn(
+            [1, model.num_expression_coeffs], dtype=torch.float32)
+
+    output = model(betas=betas, expression=expression,
+                   return_verts=True)
+    vertices = output.vertices.detach().cpu().numpy().squeeze()
+    joints = output.joints.detach().cpu().numpy().squeeze()
+
+    print('Vertices shape =', vertices.shape)
+    print('Joints shape =', joints.shape)
+
+    if plotting_module == 'pyrender':
+        import pyrender
+        import trimesh
+        vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
+        tri_mesh = trimesh.Trimesh(vertices, model.faces,
+                                   vertex_colors=vertex_colors)
+
+        mesh = pyrender.Mesh.from_trimesh(tri_mesh)
+
+        scene = pyrender.Scene()
+        scene.add(mesh)
+
+        if plot_joints:
+            sm = trimesh.creation.uv_sphere(radius=0.005)
+            sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
+            tfs = np.tile(np.eye(4), (len(joints), 1, 1))
+            tfs[:, :3, 3] = joints
+            joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
+            scene.add(joints_pcl)
+
+        pyrender.Viewer(scene, use_raymond_lighting=True)
+    elif plotting_module == 'matplotlib':
+        from matplotlib import pyplot as plt
+        from mpl_toolkits.mplot3d import Axes3D
+        from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection='3d')
+
+        mesh = Poly3DCollection(vertices[model.faces], alpha=0.1)
+        face_color = (1.0, 1.0, 0.9)
+        edge_color = (0, 0, 0)
+        mesh.set_edgecolor(edge_color)
+        mesh.set_facecolor(face_color)
+        ax.add_collection3d(mesh)
+        ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r')
+
+        if plot_joints:
+            ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1)
+        plt.show()
+    elif plotting_module == 'open3d':
+        import open3d as o3d
+
+        mesh = o3d.geometry.TriangleMesh()
+        mesh.vertices = o3d.utility.Vector3dVector(
+            vertices)
+        mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+        mesh.compute_vertex_normals()
+        mesh.paint_uniform_color([0.3, 0.3, 0.3])
+
+        geometry = [mesh]
+        if plot_joints:
+            joints_pcl = o3d.geometry.PointCloud()
+            joints_pcl.points = o3d.utility.Vector3dVector(joints)
+            joints_pcl.paint_uniform_color([0.7, 0.3, 0.3])
+            geometry.append(joints_pcl)
+
+        o3d.visualization.draw_geometries(geometry)
+    else:
+        raise ValueError('Unknown plotting_module: {}'.format(plotting_module))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+    parser.add_argument('--model-folder', required=True, type=str,
+                        help='The path to the model folder')
+    parser.add_argument('--model-type', default='smplx', type=str,
+                        choices=['smpl', 'smplh', 'smplx', 'mano', 'flame'],
+                        help='The type of model to load')
+    parser.add_argument('--gender', type=str, default='neutral',
+                        help='The gender of the model')
+    parser.add_argument('--num-betas', default=10, type=int,
+                        dest='num_betas',
+                        help='Number of shape coefficients.')
+    parser.add_argument('--num-expression-coeffs', default=10, type=int,
+                        dest='num_expression_coeffs',
+                        help='Number of expression coefficients.')
+    parser.add_argument('--plotting-module', type=str, default='pyrender',
+                        dest='plotting_module',
+                        choices=['pyrender', 'matplotlib', 'open3d'],
+                        help='The module to use for plotting the result')
+    parser.add_argument('--ext', type=str, default='npz',
+                        help='Which extension to use for loading')
+    parser.add_argument('--plot-joints', default=False,
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='The path to the model folder')
+    parser.add_argument('--sample-shape', default=True,
+                        dest='sample_shape',
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='Sample a random shape')
+    parser.add_argument('--sample-expression', default=True,
+                        dest='sample_expression',
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='Sample a random expression')
+    parser.add_argument('--use-face-contour', default=False,
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='Compute the contour of the face')
+
+    args = parser.parse_args()
+
+    model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+    model_type = args.model_type
+    plot_joints = args.plot_joints
+    use_face_contour = args.use_face_contour
+    gender = args.gender
+    ext = args.ext
+    plotting_module = args.plotting_module
+    num_betas = args.num_betas
+    num_expression_coeffs = args.num_expression_coeffs
+    sample_shape = args.sample_shape
+    sample_expression = args.sample_expression
+
+    main(model_folder, model_type, ext=ext,
+         gender=gender, plot_joints=plot_joints,
+         num_betas=num_betas,
+         num_expression_coeffs=num_expression_coeffs,
+         sample_shape=sample_shape,
+         sample_expression=sample_expression,
+         plotting_module=plotting_module,
+         use_face_contour=use_face_contour)

common/utils/smplx/examples/demo_layers.py

@@ -0,0 +1,181 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os.path as osp
+import argparse
+
+import numpy as np
+import torch
+
+import smplx
+
+
+def main(model_folder,
+         model_type='smplx',
+         ext='npz',
+         gender='neutral',
+         plot_joints=False,
+         num_betas=10,
+         sample_shape=True,
+         sample_expression=True,
+         num_expression_coeffs=10,
+         plotting_module='pyrender',
+         use_face_contour=False):
+
+    model = smplx.build_layer(
+        model_folder, model_type=model_type,
+        gender=gender, use_face_contour=use_face_contour,
+        num_betas=num_betas,
+        num_expression_coeffs=num_expression_coeffs,
+        ext=ext)
+    print(model)
+
+    betas, expression = None, None
+    if sample_shape:
+        betas = torch.randn([1, model.num_betas], dtype=torch.float32)
+    if sample_expression:
+        expression = torch.randn(
+            [1, model.num_expression_coeffs], dtype=torch.float32)
+
+    output = model(betas=betas, expression=expression,
+                   return_verts=True)
+    vertices = output.vertices.detach().cpu().numpy().squeeze()
+    joints = output.joints.detach().cpu().numpy().squeeze()
+
+    print('Vertices shape =', vertices.shape)
+    print('Joints shape =', joints.shape)
+
+    if plotting_module == 'pyrender':
+        import pyrender
+        import trimesh
+        vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
+        tri_mesh = trimesh.Trimesh(vertices, model.faces,
+                                   vertex_colors=vertex_colors)
+
+        mesh = pyrender.Mesh.from_trimesh(tri_mesh)
+
+        scene = pyrender.Scene()
+        scene.add(mesh)
+
+        if plot_joints:
+            sm = trimesh.creation.uv_sphere(radius=0.005)
+            sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
+            tfs = np.tile(np.eye(4), (len(joints), 1, 1))
+            tfs[:, :3, 3] = joints
+            joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
+            scene.add(joints_pcl)
+
+        pyrender.Viewer(scene, use_raymond_lighting=True)
+    elif plotting_module == 'matplotlib':
+        from matplotlib import pyplot as plt
+        from mpl_toolkits.mplot3d import Axes3D
+        from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection='3d')
+
+        mesh = Poly3DCollection(vertices[model.faces], alpha=0.1)
+        face_color = (1.0, 1.0, 0.9)
+        edge_color = (0, 0, 0)
+        mesh.set_edgecolor(edge_color)
+        mesh.set_facecolor(face_color)
+        ax.add_collection3d(mesh)
+        ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r')
+
+        if plot_joints:
+            ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1)
+        plt.show()
+    elif plotting_module == 'open3d':
+        import open3d as o3d
+
+        mesh = o3d.geometry.TriangleMesh()
+        mesh.vertices = o3d.utility.Vector3dVector(
+            vertices)
+        mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+        mesh.compute_vertex_normals()
+        mesh.paint_uniform_color([0.3, 0.3, 0.3])
+
+        geometry = [mesh]
+        if plot_joints:
+            joints_pcl = o3d.geometry.PointCloud()
+            joints_pcl.points = o3d.utility.Vector3dVector(joints)
+            joints_pcl.paint_uniform_color([0.7, 0.3, 0.3])
+            geometry.append(joints_pcl)
+
+        o3d.visualization.draw_geometries(geometry)
+    else:
+        raise ValueError('Unknown plotting_module: {}'.format(plotting_module))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+    parser.add_argument('--model-folder', required=True, type=str,
+                        help='The path to the model folder')
+    parser.add_argument('--model-type', default='smplx', type=str,
+                        choices=['smpl', 'smplh', 'smplx', 'mano', 'flame'],
+                        help='The type of model to load')
+    parser.add_argument('--gender', type=str, default='neutral',
+                        help='The gender of the model')
+    parser.add_argument('--num-betas', default=10, type=int,
+                        dest='num_betas',
+                        help='Number of shape coefficients.')
+    parser.add_argument('--num-expression-coeffs', default=10, type=int,
+                        dest='num_expression_coeffs',
+                        help='Number of expression coefficients.')
+    parser.add_argument('--plotting-module', type=str, default='pyrender',
+                        dest='plotting_module',
+                        choices=['pyrender', 'matplotlib', 'open3d'],
+                        help='The module to use for plotting the result')
+    parser.add_argument('--ext', type=str, default='npz',
+                        help='Which extension to use for loading')
+    parser.add_argument('--plot-joints', default=False,
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='The path to the model folder')
+    parser.add_argument('--sample-shape', default=True,
+                        dest='sample_shape',
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='Sample a random shape')
+    parser.add_argument('--sample-expression', default=True,
+                        dest='sample_expression',
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='Sample a random expression')
+    parser.add_argument('--use-face-contour', default=False,
+                        type=lambda arg: arg.lower() in ['true', '1'],
+                        help='Compute the contour of the face')
+
+    args = parser.parse_args()
+
+    model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+    model_type = args.model_type
+    plot_joints = args.plot_joints
+    use_face_contour = args.use_face_contour
+    gender = args.gender
+    ext = args.ext
+    plotting_module = args.plotting_module
+    num_betas = args.num_betas
+    num_expression_coeffs = args.num_expression_coeffs
+    sample_shape = args.sample_shape
+    sample_expression = args.sample_expression
+
+    main(model_folder, model_type, ext=ext,
+         gender=gender, plot_joints=plot_joints,
+         num_betas=num_betas,
+         num_expression_coeffs=num_expression_coeffs,
+         sample_shape=sample_shape,
+         sample_expression=sample_expression,
+         plotting_module=plotting_module,
+         use_face_contour=use_face_contour)

common/utils/smplx/examples/vis_flame_vertices.py

@@ -0,0 +1,92 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os.path as osp
+import argparse
+import pickle
+
+import numpy as np
+import torch
+import open3d as o3d
+
+import smplx
+
+
+def main(model_folder, corr_fname, ext='npz',
+         head_color=(0.3, 0.3, 0.6),
+         gender='neutral'):
+
+    head_idxs = np.load(corr_fname)
+
+    model = smplx.create(model_folder, model_type='smplx',
+                         gender=gender,
+                         ext=ext)
+    betas = torch.zeros([1, 10], dtype=torch.float32)
+    expression = torch.zeros([1, 10], dtype=torch.float32)
+
+    output = model(betas=betas, expression=expression,
+                   return_verts=True)
+    vertices = output.vertices.detach().cpu().numpy().squeeze()
+    joints = output.joints.detach().cpu().numpy().squeeze()
+
+    print('Vertices shape =', vertices.shape)
+    print('Joints shape =', joints.shape)
+
+    mesh = o3d.geometry.TriangleMesh()
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+    mesh.compute_vertex_normals()
+
+    colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
+    colors[head_idxs] = head_color
+
+    mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
+
+    o3d.visualization.draw_geometries([mesh])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+    parser.add_argument('--model-folder', required=True, type=str,
+                        help='The path to the model folder')
+    parser.add_argument('--corr-fname', required=True, type=str,
+                        dest='corr_fname',
+                        help='Filename with the head correspondences')
+    parser.add_argument('--gender', type=str, default='neutral',
+                        help='The gender of the model')
+    parser.add_argument('--ext', type=str, default='npz',
+                        help='Which extension to use for loading')
+    parser.add_argument('--head', default='right',
+                        choices=['right', 'left'],
+                        type=str, help='Which head to plot')
+    parser.add_argument('--head-color', type=float, nargs=3, dest='head_color',
+                        default=(0.3, 0.3, 0.6),
+                        help='Color for the head vertices')
+
+    args = parser.parse_args()
+
+    model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+    corr_fname = args.corr_fname
+    gender = args.gender
+    ext = args.ext
+    head = args.head
+    head_color = args.head_color
+
+    main(model_folder, corr_fname, ext=ext,
+         head_color=head_color,
+         gender=gender
+         )

common/utils/smplx/examples/vis_mano_vertices.py

@@ -0,0 +1,99 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os.path as osp
+import argparse
+import pickle
+
+import numpy as np
+import torch
+import open3d as o3d
+
+import smplx
+
+
+def main(model_folder, corr_fname, ext='npz',
+         hand_color=(0.3, 0.3, 0.6),
+         gender='neutral', hand='right'):
+
+    with open(corr_fname, 'rb') as f:
+        idxs_data = pickle.load(f)
+        if hand == 'both':
+            hand_idxs = np.concatenate(
+                [idxs_data['left_hand'], idxs_data['right_hand']]
+            )
+        else:
+            hand_idxs = idxs_data[f'{hand}_hand']
+
+    model = smplx.create(model_folder, model_type='smplx',
+                         gender=gender,
+                         ext=ext)
+    betas = torch.zeros([1, 10], dtype=torch.float32)
+    expression = torch.zeros([1, 10], dtype=torch.float32)
+
+    output = model(betas=betas, expression=expression,
+                   return_verts=True)
+    vertices = output.vertices.detach().cpu().numpy().squeeze()
+    joints = output.joints.detach().cpu().numpy().squeeze()
+
+    print('Vertices shape =', vertices.shape)
+    print('Joints shape =', joints.shape)
+
+    mesh = o3d.geometry.TriangleMesh()
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+    mesh.compute_vertex_normals()
+
+    colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
+    colors[hand_idxs] = hand_color
+
+    mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
+
+    o3d.visualization.draw_geometries([mesh])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+    parser.add_argument('--model-folder', required=True, type=str,
+                        help='The path to the model folder')
+    parser.add_argument('--corr-fname', required=True, type=str,
+                        dest='corr_fname',
+                        help='Filename with the hand correspondences')
+    parser.add_argument('--gender', type=str, default='neutral',
+                        help='The gender of the model')
+    parser.add_argument('--ext', type=str, default='npz',
+                        help='Which extension to use for loading')
+    parser.add_argument('--hand', default='right',
+                        choices=['right', 'left', 'both'],
+                        type=str, help='Which hand to plot')
+    parser.add_argument('--hand-color', type=float, nargs=3, dest='hand_color',
+                        default=(0.3, 0.3, 0.6),
+                        help='Color for the hand vertices')
+
+    args = parser.parse_args()
+
+    model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+    corr_fname = args.corr_fname
+    gender = args.gender
+    ext = args.ext
+    hand = args.hand
+    hand_color = args.hand_color
+
+    main(model_folder, corr_fname, ext=ext,
+         hand_color=hand_color,
+         gender=gender, hand=hand
+         )

common/utils/smplx/setup.py

@@ -0,0 +1,79 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: [email protected]
+
+import io
+import os
+
+from setuptools import setup
+
+# Package meta-data.
+NAME = 'smplx'
+DESCRIPTION = 'PyTorch module for loading the SMPLX body model'
+URL = 'http://smpl-x.is.tuebingen.mpg.de'
+EMAIL = '[email protected]'
+AUTHOR = 'Vassilis Choutas'
+REQUIRES_PYTHON = '>=3.6.0'
+VERSION = '0.1.21'
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+try:
+    FileNotFoundError
+except NameError:
+    FileNotFoundError = IOError
+
+# Import the README and use it as the long-description.
+# Note: this will only work if 'README.md' is present in your MANIFEST.in file!
+try:
+    with io.open(os.path.join(here, 'README.md'), encoding='utf-8') as f:
+        long_description = '\n' + f.read()
+except FileNotFoundError:
+    long_description = DESCRIPTION
+
+# Load the package's __version__.py module as a dictionary.
+about = {}
+if not VERSION:
+    with open(os.path.join(here, NAME, '__version__.py')) as f:
+        exec(f.read(), about)
+else:
+    about['__version__'] = VERSION
+
+pyrender_reqs = ['pyrender>=0.1.23', 'trimesh>=2.37.6', 'shapely']
+matplotlib_reqs = ['matplotlib']
+open3d_reqs = ['open3d-python']
+
+setup(name=NAME,
+      version=about['__version__'],
+      description=DESCRIPTION,
+      long_description=long_description,
+      long_description_content_type='text/markdown',
+      author=AUTHOR,
+      author_email=EMAIL,
+      python_requires=REQUIRES_PYTHON,
+      url=URL,
+      install_requires=[
+          'numpy>=1.16.2',
+          'torch>=1.0.1.post2',
+          'torchgeometry>=0.1.2'
+      ],
+      extras_require={
+          'pyrender': pyrender_reqs,
+          'open3d': open3d_reqs,
+          'matplotlib': matplotlib_reqs,
+          'all': pyrender_reqs + matplotlib_reqs + open3d_reqs
+      },
+      packages=['smplx', 'tools'])

common/utils/smplx/smplx/__init__.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from .body_models import (
+    create,
+    SMPL,
+    SMPLH,
+    SMPLX,
+    MANO,
+    FLAME,
+    build_layer,
+    SMPLLayer,
+    SMPLHLayer,
+    SMPLXLayer,
+    MANOLayer,
+    FLAMELayer,
+)

common/utils/smplx/smplx/body_models.py

@@ -0,0 +1,2331 @@
+#  -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from typing import Optional, Dict, Union
+import os
+import os.path as osp
+
+import pickle
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from .lbs import (
+    lbs, vertices2landmarks, find_dynamic_lmk_idx_and_bcoords)
+
+from .vertex_ids import vertex_ids as VERTEX_IDS
+from .utils import (
+    Struct, to_np, to_tensor, Tensor, Array,
+    SMPLOutput,
+    SMPLHOutput,
+    SMPLXOutput,
+    MANOOutput,
+    FLAMEOutput,
+    find_joint_kin_chain)
+from .vertex_joint_selector import VertexJointSelector
+from config import cfg
+
+class SMPL(nn.Module):
+
+    NUM_JOINTS = 23
+    NUM_BODY_JOINTS = 23
+    SHAPE_SPACE_DIM = 300
+
+    def __init__(
+        self, model_path: str,
+        data_struct: Optional[Struct] = None,
+        create_betas: bool = True,
+        betas: Optional[Tensor] = None,
+        num_betas: int = 10,
+        create_global_orient: bool = True,
+        global_orient: Optional[Tensor] = None,
+        create_body_pose: bool = True,
+        body_pose: Optional[Tensor] = None,
+        create_transl: bool = True,
+        transl: Optional[Tensor] = None,
+        dtype=torch.float32,
+        batch_size: int = 1,
+        joint_mapper=None,
+        gender: str = 'neutral',
+        vertex_ids: Dict[str, int] = None,
+        v_template: Optional[Union[Tensor, Array]] = None,
+        **kwargs
+    ) -> None:
+        ''' SMPL model constructor
+
+            Parameters
+            ----------
+            model_path: str
+                The path to the folder or to the file where the model
+                parameters are stored
+            data_struct: Strct
+                A struct object. If given, then the parameters of the model are
+                read from the object. Otherwise, the model tries to read the
+                parameters from the given `model_path`. (default = None)
+            create_global_orient: bool, optional
+                Flag for creating a member variable for the global orientation
+                of the body. (default = True)
+            global_orient: torch.tensor, optional, Bx3
+                The default value for the global orientation variable.
+                (default = None)
+            create_body_pose: bool, optional
+                Flag for creating a member variable for the pose of the body.
+                (default = True)
+            body_pose: torch.tensor, optional, Bx(Body Joints * 3)
+                The default value for the body pose variable.
+                (default = None)
+            num_betas: int, optional
+                Number of shape components to use
+                (default = 10).
+            create_betas: bool, optional
+                Flag for creating a member variable for the shape space
+                (default = True).
+            betas: torch.tensor, optional, Bx10
+                The default value for the shape member variable.
+                (default = None)
+            create_transl: bool, optional
+                Flag for creating a member variable for the translation
+                of the body. (default = True)
+            transl: torch.tensor, optional, Bx3
+                The default value for the transl variable.
+                (default = None)
+            dtype: torch.dtype, optional
+                The data type for the created variables
+            batch_size: int, optional
+                The batch size used for creating the member variables
+            joint_mapper: object, optional
+                An object that re-maps the joints. Useful if one wants to
+                re-order the SMPL joints to some other convention (e.g. MSCOCO)
+                (default = None)
+            gender: str, optional
+                Which gender to load
+            vertex_ids: dict, optional
+                A dictionary containing the indices of the extra vertices that
+                will be selected
+        '''
+
+        self.gender = gender
+
+        if data_struct is None:
+            if osp.isdir(model_path):
+                model_fn = 'SMPL_{}.{ext}'.format(gender.upper(), ext='pkl')
+                smpl_path = os.path.join(model_path, model_fn)
+            else:
+                smpl_path = model_path
+            assert osp.exists(smpl_path), 'Path {} does not exist!'.format(
+                smpl_path)
+
+            with open(smpl_path, 'rb') as smpl_file:
+                data_struct = Struct(**pickle.load(smpl_file,
+                                                   encoding='latin1'))
+
+        super(SMPL, self).__init__()
+        self.batch_size = batch_size
+        shapedirs = data_struct.shapedirs
+        if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM):
+            print(f'WARNING: You are using a {self.name()} model, with only'
+                  ' 10 shape coefficients.')
+            num_betas = min(num_betas, 10)
+        else:
+            num_betas = min(num_betas, self.SHAPE_SPACE_DIM)
+
+        self._num_betas = num_betas
+        shapedirs = shapedirs[:, :, :num_betas]
+        # The shape components
+        self.register_buffer(
+            'shapedirs',
+            to_tensor(to_np(shapedirs), dtype=dtype))
+
+        if vertex_ids is None:
+            # SMPL and SMPL-H share the same topology, so any extra joints can
+            # be drawn from the same place
+            vertex_ids = VERTEX_IDS['smplh']
+
+        self.dtype = dtype
+
+        self.joint_mapper = joint_mapper
+
+        self.vertex_joint_selector = VertexJointSelector(
+            vertex_ids=vertex_ids, **kwargs)
+
+        self.faces = data_struct.f
+        self.register_buffer('faces_tensor',
+                             to_tensor(to_np(self.faces, dtype=np.int64),
+                                       dtype=torch.long))
+
+        if create_betas:
+            if betas is None:
+                default_betas = torch.zeros(
+                    [batch_size, self.num_betas], dtype=dtype)
+            else:
+                if torch.is_tensor(betas):
+                    default_betas = betas.clone().detach()
+                else:
+                    default_betas = torch.tensor(betas, dtype=dtype)
+
+            self.register_parameter(
+                'betas', nn.Parameter(default_betas, requires_grad=True))
+
+        # The tensor that contains the global rotation of the model
+        # It is separated from the pose of the joints in case we wish to
+        # optimize only over one of them
+        if create_global_orient:
+            if global_orient is None:
+                default_global_orient = torch.zeros(
+                    [batch_size, 3], dtype=dtype)
+            else:
+                if torch.is_tensor(global_orient):
+                    default_global_orient = global_orient.clone().detach()
+                else:
+                    default_global_orient = torch.tensor(
+                        global_orient, dtype=dtype)
+
+            global_orient = nn.Parameter(default_global_orient,
+                                         requires_grad=True)
+            self.register_parameter('global_orient', global_orient)
+
+        if create_body_pose:
+            if body_pose is None:
+                default_body_pose = torch.zeros(
+                    [batch_size, self.NUM_BODY_JOINTS * 3], dtype=dtype)
+            else:
+                if torch.is_tensor(body_pose):
+                    default_body_pose = body_pose.clone().detach()
+                else:
+                    default_body_pose = torch.tensor(body_pose,
+                                                     dtype=dtype)
+            self.register_parameter(
+                'body_pose',
+                nn.Parameter(default_body_pose, requires_grad=True))
+
+        if create_transl:
+            if transl is None:
+                default_transl = torch.zeros([batch_size, 3],
+                                             dtype=dtype,
+                                             requires_grad=True)
+            else:
+                default_transl = torch.tensor(transl, dtype=dtype)
+            self.register_parameter(
+                'transl', nn.Parameter(default_transl, requires_grad=True))
+
+        if v_template is None:
+            v_template = data_struct.v_template
+        if not torch.is_tensor(v_template):
+            v_template = to_tensor(to_np(v_template), dtype=dtype)
+        # The vertices of the template model
+        self.register_buffer('v_template', v_template)
+
+        j_regressor = to_tensor(to_np(
+            data_struct.J_regressor), dtype=dtype)
+        self.register_buffer('J_regressor', j_regressor)
+
+        # Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*3 x 207
+        num_pose_basis = data_struct.posedirs.shape[-1]
+        # 207 x 20670
+        posedirs = np.reshape(data_struct.posedirs, [-1, num_pose_basis]).T
+        self.register_buffer('posedirs',
+                             to_tensor(to_np(posedirs), dtype=dtype))
+
+        # indices of parents for each joints
+        parents = to_tensor(to_np(data_struct.kintree_table[0])).long()
+        parents[0] = -1
+        self.register_buffer('parents', parents)
+
+        self.register_buffer(
+            'lbs_weights', to_tensor(to_np(data_struct.weights), dtype=dtype))
+
+    @property
+    def num_betas(self):
+        return self._num_betas
+
+    @property
+    def num_expression_coeffs(self):
+        return 0
+
+    def create_mean_pose(self, data_struct) -> Tensor:
+        pass
+
+    def name(self) -> str:
+        return 'SMPL'
+
+    @torch.no_grad()
+    def reset_params(self, **params_dict) -> None:
+        for param_name, param in self.named_parameters():
+            if param_name in params_dict:
+                param[:] = torch.tensor(params_dict[param_name])
+            else:
+                param.fill_(0)
+
+    def get_num_verts(self) -> int:
+        return self.v_template.shape[0]
+
+    def get_num_faces(self) -> int:
+        return self.faces.shape[0]
+
+    def extra_repr(self) -> str:
+        msg = [
+            f'Gender: {self.gender.upper()}',
+            f'Number of joints: {self.J_regressor.shape[0]}',
+            f'Betas: {self.num_betas}',
+        ]
+        return '\n'.join(msg)
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        body_pose: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        return_verts=True,
+        return_full_pose: bool = False,
+        pose2rot: bool = True,
+        **kwargs
+    ) -> SMPLOutput:
+        ''' Forward pass for the SMPL model
+
+            Parameters
+            ----------
+            global_orient: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable and use it as the global
+                rotation of the body. Useful if someone wishes to predicts this
+                with an external model. (default=None)
+            betas: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `betas` and use it
+                instead. For example, it can used if shape parameters
+                `betas` are predicted from some external model.
+                (default=None)
+            body_pose: torch.tensor, optional, shape Bx(J*3)
+                If given, ignore the member variable `body_pose` and use it
+                instead. For example, it can used if someone predicts the
+                pose of the body joints are predicted from some external model.
+                It should be a tensor that contains joint rotations in
+                axis-angle format. (default=None)
+            transl: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `transl` and use it
+                instead. For example, it can used if the translation
+                `transl` is predicted from some external model.
+                (default=None)
+            return_verts: bool, optional
+                Return the vertices. (default=True)
+            return_full_pose: bool, optional
+                Returns the full axis-angle pose vector (default=False)
+
+            Returns
+            -------
+        '''
+        # If no shape and pose parameters are passed along, then use the
+        # ones from the module
+        global_orient = (global_orient if global_orient is not None else
+                         self.global_orient)
+        body_pose = body_pose if body_pose is not None else self.body_pose
+        betas = betas if betas is not None else self.betas
+
+        apply_trans = transl is not None or hasattr(self, 'transl')
+        if transl is None and hasattr(self, 'transl'):
+            transl = self.transl
+
+        full_pose = torch.cat([global_orient, body_pose], dim=1)
+
+        batch_size = max(betas.shape[0], global_orient.shape[0],
+                         body_pose.shape[0])
+
+        if betas.shape[0] != batch_size:
+            num_repeats = int(batch_size / betas.shape[0])
+            betas = betas.expand(num_repeats, -1)
+
+        vertices, joints = lbs(betas, full_pose, self.v_template,
+                               self.shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=pose2rot)
+
+        joints = self.vertex_joint_selector(vertices, joints)
+        # Map the joints to the current dataset
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints)
+
+        if apply_trans:
+            joints += transl.unsqueeze(dim=1)
+            vertices += transl.unsqueeze(dim=1)
+
+        output = SMPLOutput(vertices=vertices if return_verts else None,
+                            global_orient=global_orient,
+                            body_pose=body_pose,
+                            joints=joints,
+                            betas=betas,
+                            full_pose=full_pose if return_full_pose else None)
+
+        return output
+
+
+class SMPLLayer(SMPL):
+    def __init__(
+        self,
+        *args,
+        **kwargs
+    ) -> None:
+        # Just create a SMPL module without any member variables
+        super(SMPLLayer, self).__init__(
+            create_body_pose=False,
+            create_betas=False,
+            create_global_orient=False,
+            create_transl=False,
+            *args,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        body_pose: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        return_verts=True,
+        return_full_pose: bool = False,
+        pose2rot: bool = True,
+        **kwargs
+    ) -> SMPLOutput:
+        ''' Forward pass for the SMPL model
+
+            Parameters
+            ----------
+            global_orient: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable and use it as the global
+                rotation of the body. Useful if someone wishes to predicts this
+                with an external model. (default=None)
+            betas: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `betas` and use it
+                instead. For example, it can used if shape parameters
+                `betas` are predicted from some external model.
+                (default=None)
+            body_pose: torch.tensor, optional, shape Bx(J*3)
+                If given, ignore the member variable `body_pose` and use it
+                instead. For example, it can used if someone predicts the
+                pose of the body joints are predicted from some external model.
+                It should be a tensor that contains joint rotations in
+                axis-angle format. (default=None)
+            transl: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `transl` and use it
+                instead. For example, it can used if the translation
+                `transl` is predicted from some external model.
+                (default=None)
+            return_verts: bool, optional
+                Return the vertices. (default=True)
+            return_full_pose: bool, optional
+                Returns the full axis-angle pose vector (default=False)
+
+            Returns
+            -------
+        '''
+        device, dtype = self.shapedirs.device, self.shapedirs.dtype
+        if global_orient is None:
+            batch_size = 1
+            global_orient = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 1, 1).contiguous()
+        else:
+            batch_size = global_orient.shape[0]
+        if body_pose is None:
+            body_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(
+                    batch_size, self.NUM_BODY_JOINTS, 1).contiguous()
+        if betas is None:
+            betas = torch.zeros([batch_size, self.num_betas],
+                                dtype=dtype, device=device)
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+        full_pose = torch.cat(
+            [global_orient.reshape(-1, 1, 3),
+             body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3)],
+            dim=1)
+
+        vertices, joints = lbs(betas, full_pose, self.v_template,
+                               self.shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights,
+                               pose2rot=True)
+
+        joints = self.vertex_joint_selector(vertices, joints)
+        # Map the joints to the current dataset
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints)
+
+        if transl is not None:
+            joints += transl.unsqueeze(dim=1)
+            vertices += transl.unsqueeze(dim=1)
+
+        output = SMPLOutput(vertices=vertices if return_verts else None,
+                            global_orient=global_orient,
+                            body_pose=body_pose,
+                            joints=joints,
+                            betas=betas,
+                            full_pose=full_pose if return_full_pose else None)
+
+        return output
+
+
+class SMPLH(SMPL):
+
+    # The hand joints are replaced by MANO
+    NUM_BODY_JOINTS = SMPL.NUM_JOINTS - 2
+    NUM_HAND_JOINTS = 15
+    NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS
+
+    def __init__(
+        self, model_path,
+        data_struct: Optional[Struct] = None,
+        create_left_hand_pose: bool = True,
+        left_hand_pose: Optional[Tensor] = None,
+        create_right_hand_pose: bool = True,
+        right_hand_pose: Optional[Tensor] = None,
+        use_pca: bool = True,
+        num_pca_comps: int = 6,
+        flat_hand_mean: bool = False,
+        batch_size: int = 1,
+        gender: str = 'neutral',
+        dtype=torch.float32,
+        vertex_ids=None,
+        use_compressed: bool = True,
+        ext: str = 'pkl',
+        **kwargs
+    ) -> None:
+        ''' SMPLH model constructor
+
+            Parameters
+            ----------
+            model_path: str
+                The path to the folder or to the file where the model
+                parameters are stored
+            data_struct: Strct
+                A struct object. If given, then the parameters of the model are
+                read from the object. Otherwise, the model tries to read the
+                parameters from the given `model_path`. (default = None)
+            create_left_hand_pose: bool, optional
+                Flag for creating a member variable for the pose of the left
+                hand. (default = True)
+            left_hand_pose: torch.tensor, optional, BxP
+                The default value for the left hand pose member variable.
+                (default = None)
+            create_right_hand_pose: bool, optional
+                Flag for creating a member variable for the pose of the right
+                hand. (default = True)
+            right_hand_pose: torch.tensor, optional, BxP
+                The default value for the right hand pose member variable.
+                (default = None)
+            num_pca_comps: int, optional
+                The number of PCA components to use for each hand.
+                (default = 6)
+            flat_hand_mean: bool, optional
+                If False, then the pose of the hand is initialized to False.
+            batch_size: int, optional
+                The batch size used for creating the member variables
+            gender: str, optional
+                Which gender to load
+            dtype: torch.dtype, optional
+                The data type for the created variables
+            vertex_ids: dict, optional
+                A dictionary containing the indices of the extra vertices that
+                will be selected
+        '''
+
+        self.num_pca_comps = num_pca_comps
+        # If no data structure is passed, then load the data from the given
+        # model folder
+        if data_struct is None:
+            # Load the model
+            if osp.isdir(model_path):
+                model_fn = 'SMPLH_{}.{ext}'.format(gender.upper(), ext=ext)
+                smplh_path = os.path.join(model_path, model_fn)
+            else:
+                smplh_path = model_path
+            assert osp.exists(smplh_path), 'Path {} does not exist!'.format(
+                smplh_path)
+
+            if ext == 'pkl':
+                with open(smplh_path, 'rb') as smplh_file:
+                    model_data = pickle.load(smplh_file, encoding='latin1')
+            elif ext == 'npz':
+                model_data = np.load(smplh_path, allow_pickle=True)
+            else:
+                raise ValueError('Unknown extension: {}'.format(ext))
+            data_struct = Struct(**model_data)
+
+        if vertex_ids is None:
+            vertex_ids = VERTEX_IDS['smplh']
+
+        super(SMPLH, self).__init__(
+            model_path=model_path,
+            data_struct=data_struct,
+            batch_size=batch_size, vertex_ids=vertex_ids, gender=gender,
+            use_compressed=use_compressed, dtype=dtype, ext=ext, **kwargs)
+
+        self.use_pca = use_pca
+        self.num_pca_comps = num_pca_comps
+        self.flat_hand_mean = flat_hand_mean
+
+        left_hand_components = data_struct.hands_componentsl[:num_pca_comps]
+        right_hand_components = data_struct.hands_componentsr[:num_pca_comps]
+
+        self.np_left_hand_components = left_hand_components
+        self.np_right_hand_components = right_hand_components
+        if self.use_pca:
+            self.register_buffer(
+                'left_hand_components',
+                torch.tensor(left_hand_components, dtype=dtype))
+            self.register_buffer(
+                'right_hand_components',
+                torch.tensor(right_hand_components, dtype=dtype))
+
+        if self.flat_hand_mean:
+            left_hand_mean = np.zeros_like(data_struct.hands_meanl)
+        else:
+            left_hand_mean = data_struct.hands_meanl
+
+        if self.flat_hand_mean:
+            right_hand_mean = np.zeros_like(data_struct.hands_meanr)
+        else:
+            right_hand_mean = data_struct.hands_meanr
+
+        self.register_buffer('left_hand_mean',
+                             to_tensor(left_hand_mean, dtype=self.dtype))
+        self.register_buffer('right_hand_mean',
+                             to_tensor(right_hand_mean, dtype=self.dtype))
+
+        # Create the buffers for the pose of the left hand
+        hand_pose_dim = num_pca_comps if use_pca else 3 * self.NUM_HAND_JOINTS
+        if create_left_hand_pose:
+            if left_hand_pose is None:
+                default_lhand_pose = torch.zeros([batch_size, hand_pose_dim],
+                                                 dtype=dtype)
+            else:
+                default_lhand_pose = torch.tensor(left_hand_pose, dtype=dtype)
+
+            left_hand_pose_param = nn.Parameter(default_lhand_pose,
+                                                requires_grad=True)
+            self.register_parameter('left_hand_pose',
+                                    left_hand_pose_param)
+
+        if create_right_hand_pose:
+            if right_hand_pose is None:
+                default_rhand_pose = torch.zeros([batch_size, hand_pose_dim],
+                                                 dtype=dtype)
+            else:
+                default_rhand_pose = torch.tensor(right_hand_pose, dtype=dtype)
+
+            right_hand_pose_param = nn.Parameter(default_rhand_pose,
+                                                 requires_grad=True)
+            self.register_parameter('right_hand_pose',
+                                    right_hand_pose_param)
+
+        # Create the buffer for the mean pose.
+        pose_mean_tensor = self.create_mean_pose(
+            data_struct, flat_hand_mean=flat_hand_mean)
+        if not torch.is_tensor(pose_mean_tensor):
+            pose_mean_tensor = torch.tensor(pose_mean_tensor, dtype=dtype)
+        self.register_buffer('pose_mean', pose_mean_tensor)
+
+    def create_mean_pose(self, data_struct, flat_hand_mean=False):
+        # Create the array for the mean pose. If flat_hand is false, then use
+        # the mean that is given by the data, rather than the flat open hand
+        global_orient_mean = torch.zeros([3], dtype=self.dtype)
+        body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3],
+                                     dtype=self.dtype)
+
+        pose_mean = torch.cat([global_orient_mean, body_pose_mean,
+                               self.left_hand_mean,
+                               self.right_hand_mean], dim=0)
+        return pose_mean
+
+    def name(self) -> str:
+        return 'SMPL+H'
+
+    def extra_repr(self):
+        msg = super(SMPLH, self).extra_repr()
+        msg = [msg]
+        if self.use_pca:
+            msg.append(f'Number of PCA components: {self.num_pca_comps}')
+        msg.append(f'Flat hand mean: {self.flat_hand_mean}')
+        return '\n'.join(msg)
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        body_pose: Optional[Tensor] = None,
+        left_hand_pose: Optional[Tensor] = None,
+        right_hand_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        pose2rot: bool = True,
+        **kwargs
+    ) -> SMPLHOutput:
+        '''
+        '''
+        # If no shape and pose parameters are passed along, then use the
+        # ones from the module
+        global_orient = (global_orient if global_orient is not None else
+                         self.global_orient)
+        body_pose = body_pose if body_pose is not None else self.body_pose
+        betas = betas if betas is not None else self.betas
+        left_hand_pose = (left_hand_pose if left_hand_pose is not None else
+                          self.left_hand_pose)
+        right_hand_pose = (right_hand_pose if right_hand_pose is not None else
+                           self.right_hand_pose)
+
+        apply_trans = transl is not None or hasattr(self, 'transl')
+        if transl is None:
+            if hasattr(self, 'transl'):
+                transl = self.transl
+
+        if self.use_pca:
+            left_hand_pose = torch.einsum(
+                'bi,ij->bj', [left_hand_pose, self.left_hand_components])
+            right_hand_pose = torch.einsum(
+                'bi,ij->bj', [right_hand_pose, self.right_hand_components])
+
+        full_pose = torch.cat([global_orient, body_pose,
+                               left_hand_pose,
+                               right_hand_pose], dim=1)
+        full_pose += self.pose_mean
+
+        vertices, joints = lbs(self.betas, full_pose, self.v_template,
+                               self.shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=pose2rot)
+
+        # Add any extra joints that might be needed
+        joints = self.vertex_joint_selector(vertices, joints)
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints)
+
+        if apply_trans:
+            joints += transl.unsqueeze(dim=1)
+            vertices += transl.unsqueeze(dim=1)
+
+        output = SMPLHOutput(vertices=vertices if return_verts else None,
+                             joints=joints,
+                             betas=betas,
+                             global_orient=global_orient,
+                             body_pose=body_pose,
+                             left_hand_pose=left_hand_pose,
+                             right_hand_pose=right_hand_pose,
+                             full_pose=full_pose if return_full_pose else None)
+
+        return output
+
+
+class SMPLHLayer(SMPLH):
+
+    def __init__(
+        self, *args, **kwargs
+    ) -> None:
+        ''' SMPL+H as a layer model constructor
+        '''
+        super(SMPLHLayer, self).__init__(
+            create_global_orient=False,
+            create_body_pose=False,
+            create_left_hand_pose=False,
+            create_right_hand_pose=False,
+            create_betas=False,
+            create_transl=False,
+            *args,
+            **kwargs)
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        body_pose: Optional[Tensor] = None,
+        left_hand_pose: Optional[Tensor] = None,
+        right_hand_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        pose2rot: bool = True,
+        **kwargs
+    ) -> SMPLHOutput:
+        '''
+        '''
+        device, dtype = self.shapedirs.device, self.shapedirs.dtype
+        if global_orient is None:
+            batch_size = 1
+            global_orient = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        else:
+            batch_size = global_orient.shape[0]
+        if body_pose is None:
+            body_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 21, -1).contiguous()
+        if left_hand_pose is None:
+            left_hand_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 15, -1).contiguous()
+        if right_hand_pose is None:
+            right_hand_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 15, -1).contiguous()
+        if betas is None:
+            betas = torch.zeros([batch_size, self.num_betas],
+                                dtype=dtype, device=device)
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+        # Concatenate all pose vectors
+        full_pose = torch.cat(
+            [global_orient.reshape(-1, 1, 3),
+             body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3),
+             left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3),
+             right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3)],
+            dim=1)
+
+        vertices, joints = lbs(betas, full_pose, self.v_template,
+                               self.shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=True)
+
+        # Add any extra joints that might be needed
+        joints = self.vertex_joint_selector(vertices, joints)
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints)
+
+        if transl is not None:
+            joints += transl.unsqueeze(dim=1)
+            vertices += transl.unsqueeze(dim=1)
+
+        output = SMPLHOutput(vertices=vertices if return_verts else None,
+                             joints=joints,
+                             betas=betas,
+                             global_orient=global_orient,
+                             body_pose=body_pose,
+                             left_hand_pose=left_hand_pose,
+                             right_hand_pose=right_hand_pose,
+                             full_pose=full_pose if return_full_pose else None)
+
+        return output
+
+
+class SMPLX(SMPLH):
+    '''
+    SMPL-X (SMPL eXpressive) is a unified body model, with shape parameters
+    trained jointly for the face, hands and body.
+    SMPL-X uses standard vertex based linear blend skinning with learned
+    corrective blend shapes, has N=10475 vertices and K=54 joints,
+    which includes joints for the neck, jaw, eyeballs and fingers.
+    '''
+
+    NUM_BODY_JOINTS = SMPLH.NUM_BODY_JOINTS
+    NUM_HAND_JOINTS = 15
+    NUM_FACE_JOINTS = 3
+    NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS + NUM_FACE_JOINTS
+    EXPRESSION_SPACE_DIM = 100
+    NECK_IDX = 12
+
+    def __init__(
+        self, model_path: str,
+        num_expression_coeffs: int = 10,
+        create_expression: bool = True,
+        expression: Optional[Tensor] = None,
+        create_jaw_pose: bool = True,
+        jaw_pose: Optional[Tensor] = None,
+        create_leye_pose: bool = True,
+        leye_pose: Optional[Tensor] = None,
+        create_reye_pose=True,
+        reye_pose: Optional[Tensor] = None,
+        use_face_contour: bool = False,
+        batch_size: int = 1,
+        gender: str = 'neutral',
+        dtype=torch.float32,
+        ext: str = 'npz',
+        **kwargs
+    ) -> None:
+        ''' SMPLX model constructor
+
+            Parameters
+            ----------
+            model_path: str
+                The path to the folder or to the file where the model
+                parameters are stored
+            num_expression_coeffs: int, optional
+                Number of expression components to use
+                (default = 10).
+            create_expression: bool, optional
+                Flag for creating a member variable for the expression space
+                (default = True).
+            expression: torch.tensor, optional, Bx10
+                The default value for the expression member variable.
+                (default = None)
+            create_jaw_pose: bool, optional
+                Flag for creating a member variable for the jaw pose.
+                (default = False)
+            jaw_pose: torch.tensor, optional, Bx3
+                The default value for the jaw pose variable.
+                (default = None)
+            create_leye_pose: bool, optional
+                Flag for creating a member variable for the left eye pose.
+                (default = False)
+            leye_pose: torch.tensor, optional, Bx10
+                The default value for the left eye pose variable.
+                (default = None)
+            create_reye_pose: bool, optional
+                Flag for creating a member variable for the right eye pose.
+                (default = False)
+            reye_pose: torch.tensor, optional, Bx10
+                The default value for the right eye pose variable.
+                (default = None)
+            use_face_contour: bool, optional
+                Whether to compute the keypoints that form the facial contour
+            batch_size: int, optional
+                The batch size used for creating the member variables
+            gender: str, optional
+                Which gender to load
+            dtype: torch.dtype
+                The data type for the created variables
+        '''
+
+        # Load the model
+        if osp.isdir(model_path):
+            model_fn = 'SMPLX_{}.{ext}'.format(gender.upper(), ext=ext)
+            smplx_path = os.path.join(model_path, model_fn)
+        else:
+            smplx_path = model_path
+        assert osp.exists(smplx_path), 'Path {} does not exist!'.format(smplx_path)
+        if ext == 'pkl':
+            with open(smplx_path, 'rb') as smplx_file:
+                model_data = pickle.load(smplx_file, encoding='latin1')
+        elif ext == 'npz':
+            model_data = np.load(smplx_path, allow_pickle=True)
+        else:
+            raise ValueError('Unknown extension: {}'.format(ext))
+
+        data_struct = Struct(**model_data)
+
+        super(SMPLX, self).__init__(
+            model_path=model_path,
+            data_struct=data_struct,
+            dtype=dtype,
+            batch_size=batch_size,
+            vertex_ids=VERTEX_IDS['smplx'],
+            gender=gender, ext=ext,
+            **kwargs)
+
+        lmk_faces_idx = data_struct.lmk_faces_idx
+        self.register_buffer('lmk_faces_idx',
+                             torch.tensor(lmk_faces_idx, dtype=torch.long))
+        lmk_bary_coords = data_struct.lmk_bary_coords
+        self.register_buffer('lmk_bary_coords',
+                             torch.tensor(lmk_bary_coords, dtype=dtype))
+
+        self.use_face_contour = use_face_contour
+        if self.use_face_contour:
+            dynamic_lmk_faces_idx = data_struct.dynamic_lmk_faces_idx
+            dynamic_lmk_faces_idx = torch.tensor(
+                dynamic_lmk_faces_idx,
+                dtype=torch.long)
+            self.register_buffer('dynamic_lmk_faces_idx',
+                                 dynamic_lmk_faces_idx)
+
+            dynamic_lmk_bary_coords = data_struct.dynamic_lmk_bary_coords
+            dynamic_lmk_bary_coords = torch.tensor(
+                dynamic_lmk_bary_coords, dtype=dtype)
+            self.register_buffer('dynamic_lmk_bary_coords',
+                                 dynamic_lmk_bary_coords)
+
+            neck_kin_chain = find_joint_kin_chain(self.NECK_IDX, self.parents)
+            self.register_buffer(
+                'neck_kin_chain',
+                torch.tensor(neck_kin_chain, dtype=torch.long))
+
+        if create_jaw_pose:
+            if jaw_pose is None:
+                default_jaw_pose = torch.zeros([batch_size, 3], dtype=dtype)
+            else:
+                default_jaw_pose = torch.tensor(jaw_pose, dtype=dtype)
+            jaw_pose_param = nn.Parameter(default_jaw_pose,
+                                          requires_grad=True)
+            self.register_parameter('jaw_pose', jaw_pose_param)
+
+        if create_leye_pose:
+            if leye_pose is None:
+                default_leye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+            else:
+                default_leye_pose = torch.tensor(leye_pose, dtype=dtype)
+            leye_pose_param = nn.Parameter(default_leye_pose,
+                                           requires_grad=True)
+            self.register_parameter('leye_pose', leye_pose_param)
+
+        if create_reye_pose:
+            if reye_pose is None:
+                default_reye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+            else:
+                default_reye_pose = torch.tensor(reye_pose, dtype=dtype)
+            reye_pose_param = nn.Parameter(default_reye_pose,
+                                           requires_grad=True)
+            self.register_parameter('reye_pose', reye_pose_param)
+
+        shapedirs = data_struct.shapedirs
+        if len(shapedirs.shape) < 3:
+            shapedirs = shapedirs[:, :, None]
+        if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM +
+                self.EXPRESSION_SPACE_DIM):
+            print(f'WARNING: You are using a {self.name()} model, with only'
+                  ' 10 shape and 10 expression coefficients.')
+            expr_start_idx = 10
+            expr_end_idx = 20
+            num_expression_coeffs = min(num_expression_coeffs, 10)
+        else:
+            expr_start_idx = self.SHAPE_SPACE_DIM
+            expr_end_idx = self.SHAPE_SPACE_DIM + num_expression_coeffs
+            num_expression_coeffs = min(
+                num_expression_coeffs, self.EXPRESSION_SPACE_DIM)
+
+        self._num_expression_coeffs = num_expression_coeffs
+
+        expr_dirs = shapedirs[:, :, expr_start_idx:expr_end_idx]
+        self.register_buffer(
+            'expr_dirs', to_tensor(to_np(expr_dirs), dtype=dtype))
+
+        if create_expression:
+            if expression is None:
+                default_expression = torch.zeros(
+                    [batch_size, self.num_expression_coeffs], dtype=dtype)
+            else:
+                default_expression = torch.tensor(expression, dtype=dtype)
+            expression_param = nn.Parameter(default_expression,
+                                            requires_grad=True)
+            self.register_parameter('expression', expression_param)
+
+    def name(self) -> str:
+        return 'SMPL-X'
+
+    @property
+    def num_expression_coeffs(self):
+        return self._num_expression_coeffs
+
+    def create_mean_pose(self, data_struct, flat_hand_mean=False):
+        # Create the array for the mean pose. If flat_hand is false, then use
+        # the mean that is given by the data, rather than the flat open hand
+        global_orient_mean = torch.zeros([3], dtype=self.dtype)
+        body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3],
+                                     dtype=self.dtype)
+        jaw_pose_mean = torch.zeros([3], dtype=self.dtype)
+        leye_pose_mean = torch.zeros([3], dtype=self.dtype)
+        reye_pose_mean = torch.zeros([3], dtype=self.dtype)
+
+        pose_mean = np.concatenate([global_orient_mean, body_pose_mean,
+                                    jaw_pose_mean,
+                                    leye_pose_mean, reye_pose_mean,
+                                    self.left_hand_mean, self.right_hand_mean],
+                                   axis=0)
+
+        return pose_mean
+
+    def extra_repr(self):
+        msg = super(SMPLX, self).extra_repr()
+        msg = [
+            msg,
+            f'Number of Expression Coefficients: {self.num_expression_coeffs}'
+        ]
+        return '\n'.join(msg)
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        body_pose: Optional[Tensor] = None,
+        left_hand_pose: Optional[Tensor] = None,
+        right_hand_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        expression: Optional[Tensor] = None,
+        jaw_pose: Optional[Tensor] = None,
+        leye_pose: Optional[Tensor] = None,
+        reye_pose: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        pose2rot: bool = True,
+        **kwargs
+    ) -> SMPLXOutput:
+        '''
+        Forward pass for the SMPLX model
+
+            Parameters
+            ----------
+            global_orient: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable and use it as the global
+                rotation of the body. Useful if someone wishes to predicts this
+                with an external model. (default=None)
+            betas: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `betas` and use it
+                instead. For example, it can used if shape parameters
+                `betas` are predicted from some external model.
+                (default=None)
+            expression: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `expression` and use it
+                instead. For example, it can used if expression parameters
+                `expression` are predicted from some external model.
+            body_pose: torch.tensor, optional, shape Bx(J*3)
+                If given, ignore the member variable `body_pose` and use it
+                instead. For example, it can used if someone predicts the
+                pose of the body joints are predicted from some external model.
+                It should be a tensor that contains joint rotations in
+                axis-angle format. (default=None)
+            left_hand_pose: torch.tensor, optional, shape BxP
+                If given, ignore the member variable `left_hand_pose` and
+                use this instead. It should either contain PCA coefficients or
+                joint rotations in axis-angle format.
+            right_hand_pose: torch.tensor, optional, shape BxP
+                If given, ignore the member variable `right_hand_pose` and
+                use this instead. It should either contain PCA coefficients or
+                joint rotations in axis-angle format.
+            jaw_pose: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `jaw_pose` and
+                use this instead. It should either joint rotations in
+                axis-angle format.
+            transl: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `transl` and use it
+                instead. For example, it can used if the translation
+                `transl` is predicted from some external model.
+                (default=None)
+            return_verts: bool, optional
+                Return the vertices. (default=True)
+            return_full_pose: bool, optional
+                Returns the full axis-angle pose vector (default=False)
+
+            Returns
+            -------
+                output: ModelOutput
+                A named tuple of type `ModelOutput`
+        '''
+
+        # If no shape and pose parameters are passed along, then use the
+        # ones from the module
+        global_orient = (global_orient if global_orient is not None else
+                         self.global_orient)
+        body_pose = body_pose if body_pose is not None else self.body_pose
+        betas = betas if betas is not None else self.betas
+
+        left_hand_pose = (left_hand_pose if left_hand_pose is not None else
+                          self.left_hand_pose)
+        right_hand_pose = (right_hand_pose if right_hand_pose is not None else
+                           self.right_hand_pose)
+        jaw_pose = jaw_pose if jaw_pose is not None else self.jaw_pose
+        leye_pose = leye_pose if leye_pose is not None else self.leye_pose
+        reye_pose = reye_pose if reye_pose is not None else self.reye_pose
+        expression = expression if expression is not None else self.expression
+
+        apply_trans = transl is not None or hasattr(self, 'transl')
+        if transl is None:
+            if hasattr(self, 'transl'):
+                transl = self.transl
+
+        if self.use_pca:
+            left_hand_pose = torch.einsum(
+                'bi,ij->bj', [left_hand_pose, self.left_hand_components])
+            right_hand_pose = torch.einsum(
+                'bi,ij->bj', [right_hand_pose, self.right_hand_components])
+
+        full_pose = torch.cat([global_orient, body_pose,
+                               jaw_pose, leye_pose, reye_pose,
+                               left_hand_pose,
+                               right_hand_pose], dim=1)
+
+        # Add the mean pose of the model. Does not affect the body, only the
+        # hands when flat_hand_mean == False
+        full_pose += self.pose_mean
+
+        batch_size = max(betas.shape[0], global_orient.shape[0],
+                         body_pose.shape[0])
+        # Concatenate the shape and expression coefficients
+        scale = int(batch_size / betas.shape[0])
+        if scale > 1:
+            betas = betas.expand(scale, -1)
+        shape_components = torch.cat([betas, expression], dim=-1)
+
+        shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+        vertices, joints = lbs(shape_components, full_pose, self.v_template,
+                               shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=pose2rot,
+                               )
+
+        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(
+            dim=0).expand(batch_size, -1).contiguous()
+        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(
+            self.batch_size, 1, 1)
+        if self.use_face_contour:
+            lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+                vertices, full_pose, self.dynamic_lmk_faces_idx,
+                self.dynamic_lmk_bary_coords,
+                self.neck_kin_chain,
+                pose2rot=True,
+            )
+            dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+
+            lmk_faces_idx = torch.cat([lmk_faces_idx,
+                                       dyn_lmk_faces_idx], 1)
+            lmk_bary_coords = torch.cat(
+                [lmk_bary_coords.expand(batch_size, -1, -1),
+                 dyn_lmk_bary_coords], 1)
+
+        landmarks = vertices2landmarks(vertices, self.faces_tensor,
+                                       lmk_faces_idx,
+                                       lmk_bary_coords)
+
+        # Add any extra joints that might be needed
+        joints = self.vertex_joint_selector(vertices, joints)
+        # Add the landmarks to the joints
+        joints = torch.cat([joints, landmarks], dim=1)
+        # Map the joints to the current dataset
+
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+        if apply_trans:
+            joints += transl.unsqueeze(dim=1)
+            vertices += transl.unsqueeze(dim=1)
+
+        output = SMPLXOutput(vertices=vertices if return_verts else None,
+                             joints=joints,
+                             betas=betas,
+                             expression=expression,
+                             global_orient=global_orient,
+                             body_pose=body_pose,
+                             left_hand_pose=left_hand_pose,
+                             right_hand_pose=right_hand_pose,
+                             jaw_pose=jaw_pose,
+                             full_pose=full_pose if return_full_pose else None)
+        return output
+
+
+class SMPLXLayer(SMPLX):
+    def __init__(
+        self,
+        *args,
+        **kwargs
+    ) -> None:
+        # Just create a SMPLX module without any member variables
+        super(SMPLXLayer, self).__init__(
+            create_global_orient=False,
+            create_body_pose=False,
+            create_left_hand_pose=False,
+            create_right_hand_pose=False,
+            create_jaw_pose=False,
+            create_leye_pose=False,
+            create_reye_pose=False,
+            create_betas=False,
+            create_expression=False,
+            create_transl=False,
+            *args, **kwargs,
+        )
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        body_pose: Optional[Tensor] = None,
+        left_hand_pose: Optional[Tensor] = None,
+        right_hand_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        expression: Optional[Tensor] = None,
+        jaw_pose: Optional[Tensor] = None,
+        leye_pose: Optional[Tensor] = None,
+        reye_pose: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        **kwargs
+    ) -> SMPLXOutput:
+        '''
+        Forward pass for the SMPLX model
+
+            Parameters
+            ----------
+            global_orient: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable and use it as the global
+                rotation of the body. Useful if someone wishes to predicts this
+                with an external model. (default=None)
+            betas: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `betas` and use it
+                instead. For example, it can used if shape parameters
+                `betas` are predicted from some external model.
+                (default=None)
+            expression: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `expression` and use it
+                instead. For example, it can used if expression parameters
+                `expression` are predicted from some external model.
+            body_pose: torch.tensor, optional, shape Bx(J*3)
+                If given, ignore the member variable `body_pose` and use it
+                instead. For example, it can used if someone predicts the
+                pose of the body joints are predicted from some external model.
+                It should be a tensor that contains joint rotations in
+                axis-angle format. (default=None)
+            left_hand_pose: torch.tensor, optional, shape BxP
+                If given, ignore the member variable `left_hand_pose` and
+                use this instead. It should either contain PCA coefficients or
+                joint rotations in axis-angle format.
+            right_hand_pose: torch.tensor, optional, shape BxP
+                If given, ignore the member variable `right_hand_pose` and
+                use this instead. It should either contain PCA coefficients or
+                joint rotations in axis-angle format.
+            jaw_pose: torch.tensor, optional, shape Bx3x3
+                If given, ignore the member variable `jaw_pose` and
+                use this instead. It should either joint rotations in
+                axis-angle format.
+            transl: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `transl` and use it
+                instead. For example, it can used if the translation
+                `transl` is predicted from some external model.
+                (default=None)
+            return_verts: bool, optional
+                Return the vertices. (default=True)
+            return_full_pose: bool, optional
+                Returns the full pose vector (default=False)
+            Returns
+            -------
+                output: ModelOutput
+                A data class that contains the posed vertices and joints
+        '''
+        device, dtype = self.shapedirs.device, self.shapedirs.dtype
+
+        if global_orient is None:
+            batch_size = 1
+            global_orient = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        else:
+            batch_size = global_orient.shape[0]
+        if body_pose is None:
+            body_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(
+                    batch_size, self.NUM_BODY_JOINTS, -1).contiguous()
+        if left_hand_pose is None:
+            left_hand_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 15, -1).contiguous()
+        if right_hand_pose is None:
+            right_hand_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 15, -1).contiguous()
+        if jaw_pose is None:
+            jaw_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        if leye_pose is None:
+            leye_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        if reye_pose is None:
+            reye_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        if expression is None:
+            expression = torch.zeros([batch_size, self.num_expression_coeffs],
+                                     dtype=dtype, device=device)
+        if betas is None:
+            betas = torch.zeros([batch_size, self.num_betas],
+                                dtype=dtype, device=device)
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+        # Concatenate all pose vectors
+        full_pose = torch.cat(
+            [global_orient.reshape(-1, 1, 3),
+             body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3),
+             jaw_pose.reshape(-1, 1, 3),
+             leye_pose.reshape(-1, 1, 3),
+             reye_pose.reshape(-1, 1, 3),
+             left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3),
+             right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3)],
+            dim=1)
+        shape_components = torch.cat([betas, expression], dim=-1)
+
+        shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+        vertices, joints = lbs(shape_components, full_pose, self.v_template,
+                               shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=True)
+
+        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(
+            dim=0).expand(batch_size, -1).contiguous()
+        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(
+            self.batch_size, 1, 1)
+        if self.use_face_contour:
+            lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+                vertices, full_pose,
+                self.dynamic_lmk_faces_idx,
+                self.dynamic_lmk_bary_coords,
+                self.neck_kin_chain,
+                pose2rot=False,
+            )
+            dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+
+            lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1)
+            lmk_bary_coords = torch.cat(
+                [lmk_bary_coords.expand(batch_size, -1, -1),
+                 dyn_lmk_bary_coords], 1)
+
+        landmarks = vertices2landmarks(vertices, self.faces_tensor,
+                                       lmk_faces_idx,
+                                       lmk_bary_coords)
+
+        # Add any extra joints that might be needed
+        joints = self.vertex_joint_selector(vertices, joints)
+        # Add the landmarks to the joints
+        joints = torch.cat([joints, landmarks], dim=1)
+        # Map the joints to the current dataset
+
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+        if transl is not None:
+            joints += transl.unsqueeze(dim=1)
+            vertices += transl.unsqueeze(dim=1)
+
+        output = SMPLXOutput(vertices=vertices if return_verts else None,
+                             joints=joints,
+                             betas=betas,
+                             expression=expression,
+                             global_orient=global_orient,
+                             body_pose=body_pose,
+                             left_hand_pose=left_hand_pose,
+                             right_hand_pose=right_hand_pose,
+                             jaw_pose=jaw_pose,
+                             transl=transl,
+                             full_pose=full_pose if return_full_pose else None)
+        return output
+
+
+class MANO(SMPL):
+    # The hand joints are replaced by MANO
+    NUM_BODY_JOINTS = 1
+    NUM_HAND_JOINTS = 15
+    NUM_JOINTS = NUM_BODY_JOINTS + NUM_HAND_JOINTS
+
+    def __init__(
+        self,
+        model_path: str,
+        is_rhand: bool = True,
+        data_struct: Optional[Struct] = None,
+        create_hand_pose: bool = True,
+        hand_pose: Optional[Tensor] = None,
+        use_pca: bool = True,
+        num_pca_comps: int = 6,
+        flat_hand_mean: bool = False,
+        batch_size: int = 1,
+        dtype=torch.float32,
+        vertex_ids=None,
+        use_compressed: bool = True,
+        ext: str = 'pkl',
+        **kwargs
+    ) -> None:
+        ''' MANO model constructor
+
+            Parameters
+            ----------
+            model_path: str
+                The path to the folder or to the file where the model
+                parameters are stored
+            data_struct: Strct
+                A struct object. If given, then the parameters of the model are
+                read from the object. Otherwise, the model tries to read the
+                parameters from the given `model_path`. (default = None)
+            create_hand_pose: bool, optional
+                Flag for creating a member variable for the pose of the right
+                hand. (default = True)
+            hand_pose: torch.tensor, optional, BxP
+                The default value for the right hand pose member variable.
+                (default = None)
+            num_pca_comps: int, optional
+                The number of PCA components to use for each hand.
+                (default = 6)
+            flat_hand_mean: bool, optional
+                If False, then the pose of the hand is initialized to False.
+            batch_size: int, optional
+                The batch size used for creating the member variables
+            dtype: torch.dtype, optional
+                The data type for the created variables
+            vertex_ids: dict, optional
+                A dictionary containing the indices of the extra vertices that
+                will be selected
+        '''
+
+        self.num_pca_comps = num_pca_comps
+        self.is_rhand = is_rhand
+        # If no data structure is passed, then load the data from the given
+        # model folder
+        if data_struct is None:
+            # Load the model
+            if osp.isdir(model_path):
+                model_fn = 'MANO_{}.{ext}'.format(
+                    'RIGHT' if is_rhand else 'LEFT', ext=ext)
+                mano_path = os.path.join(model_path, model_fn)
+            else:
+                mano_path = model_path
+                self.is_rhand = True if 'RIGHT' in os.path.basename(
+                    model_path) else False
+            assert osp.exists(mano_path), 'Path {} does not exist!'.format(
+                mano_path)
+
+            if ext == 'pkl':
+                with open(mano_path, 'rb') as mano_file:
+                    model_data = pickle.load(mano_file, encoding='latin1')
+            elif ext == 'npz':
+                model_data = np.load(mano_path, allow_pickle=True)
+            else:
+                raise ValueError('Unknown extension: {}'.format(ext))
+            data_struct = Struct(**model_data)
+
+        if vertex_ids is None:
+            vertex_ids = VERTEX_IDS['smplh']
+
+        super(MANO, self).__init__(
+            model_path=model_path, data_struct=data_struct,
+            batch_size=batch_size, vertex_ids=vertex_ids,
+            use_compressed=use_compressed, dtype=dtype, ext=ext, **kwargs)
+
+        # add only MANO tips to the extra joints
+        self.vertex_joint_selector.extra_joints_idxs = to_tensor(
+            list(VERTEX_IDS['mano'].values()), dtype=torch.long)
+
+        self.use_pca = use_pca
+        self.num_pca_comps = num_pca_comps
+        if self.num_pca_comps == 45:
+            self.use_pca = False
+        self.flat_hand_mean = flat_hand_mean
+
+        hand_components = data_struct.hands_components[:num_pca_comps]
+
+        self.np_hand_components = hand_components
+
+        if self.use_pca:
+            self.register_buffer(
+                'hand_components',
+                torch.tensor(hand_components, dtype=dtype))
+
+        if self.flat_hand_mean:
+            hand_mean = np.zeros_like(data_struct.hands_mean)
+        else:
+            hand_mean = data_struct.hands_mean
+
+        self.register_buffer('hand_mean',
+                             to_tensor(hand_mean, dtype=self.dtype))
+
+        # Create the buffers for the pose of the left hand
+        hand_pose_dim = num_pca_comps if use_pca else 3 * self.NUM_HAND_JOINTS
+        if create_hand_pose:
+            if hand_pose is None:
+                default_hand_pose = torch.zeros([batch_size, hand_pose_dim],
+                                                dtype=dtype)
+            else:
+                default_hand_pose = torch.tensor(hand_pose, dtype=dtype)
+
+            hand_pose_param = nn.Parameter(default_hand_pose,
+                                           requires_grad=True)
+            self.register_parameter('hand_pose',
+                                    hand_pose_param)
+
+        # Create the buffer for the mean pose.
+        pose_mean = self.create_mean_pose(
+            data_struct, flat_hand_mean=flat_hand_mean)
+        pose_mean_tensor = pose_mean.clone().to(dtype)
+        # pose_mean_tensor = torch.tensor(pose_mean, dtype=dtype)
+        self.register_buffer('pose_mean', pose_mean_tensor)
+
+    def name(self) -> str:
+        return 'MANO'
+
+    def create_mean_pose(self, data_struct, flat_hand_mean=False):
+        # Create the array for the mean pose. If flat_hand is false, then use
+        # the mean that is given by the data, rather than the flat open hand
+        global_orient_mean = torch.zeros([3], dtype=self.dtype)
+        pose_mean = torch.cat([global_orient_mean, self.hand_mean], dim=0)
+        return pose_mean
+
+    def extra_repr(self):
+        msg = [super(MANO, self).extra_repr()]
+        if self.use_pca:
+            msg.append(f'Number of PCA components: {self.num_pca_comps}')
+        msg.append(f'Flat hand mean: {self.flat_hand_mean}')
+        return '\n'.join(msg)
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        hand_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        **kwargs
+    ) -> MANOOutput:
+        ''' Forward pass for the MANO model
+        '''
+        # If no shape and pose parameters are passed along, then use the
+        # ones from the module
+        global_orient = (global_orient if global_orient is not None else
+                         self.global_orient)
+        betas = betas if betas is not None else self.betas
+        hand_pose = (hand_pose if hand_pose is not None else
+                     self.hand_pose)
+
+        apply_trans = transl is not None or hasattr(self, 'transl')
+        if transl is None:
+            if hasattr(self, 'transl'):
+                transl = self.transl
+
+        if self.use_pca:
+            hand_pose = torch.einsum(
+                'bi,ij->bj', [hand_pose, self.hand_components])
+
+        full_pose = torch.cat([global_orient, hand_pose], dim=1)
+        full_pose += self.pose_mean
+
+        vertices, joints = lbs(betas, full_pose, self.v_template,
+                               self.shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=True,
+                               )
+
+        # # Add pre-selected extra joints that might be needed
+        # joints = self.vertex_joint_selector(vertices, joints)
+
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints)
+
+        if apply_trans:
+            joints = joints + transl.unsqueeze(dim=1)
+            vertices = vertices + transl.unsqueeze(dim=1)
+
+        output = MANOOutput(vertices=vertices if return_verts else None,
+                            joints=joints if return_verts else None,
+                            betas=betas,
+                            global_orient=global_orient,
+                            hand_pose=hand_pose,
+                            full_pose=full_pose if return_full_pose else None)
+
+        return output
+
+
+class MANOLayer(MANO):
+    def __init__(self, *args, **kwargs) -> None:
+        ''' MANO as a layer model constructor
+        '''
+        super(MANOLayer, self).__init__(
+            create_global_orient=False,
+            create_hand_pose=False,
+            create_betas=False,
+            create_transl=False,
+            *args, **kwargs)
+
+    def name(self) -> str:
+        return 'MANO'
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        hand_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        **kwargs
+    ) -> MANOOutput:
+        ''' Forward pass for the MANO model
+        '''
+        device, dtype = self.shapedirs.device, self.shapedirs.dtype
+        if global_orient is None:
+            batch_size = 1
+            global_orient = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        else:
+            batch_size = global_orient.shape[0]
+        if hand_pose is None:
+            hand_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 15, -1).contiguous()
+        if betas is None:
+            betas = torch.zeros(
+                [batch_size, self.num_betas], dtype=dtype, device=device)
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+        full_pose = torch.cat([global_orient, hand_pose], dim=1)
+        vertices, joints = lbs(betas, full_pose, self.v_template,
+                               self.shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=True)
+
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints)
+
+        if transl is not None:
+            joints = joints + transl.unsqueeze(dim=1)
+            vertices = vertices + transl.unsqueeze(dim=1)
+
+        output = MANOOutput(
+            vertices=vertices if return_verts else None,
+            joints=joints if return_verts else None,
+            betas=betas,
+            global_orient=global_orient,
+            hand_pose=hand_pose,
+            full_pose=full_pose if return_full_pose else None)
+
+        return output
+
+
+class FLAME(SMPL):
+    NUM_JOINTS = 5
+    SHAPE_SPACE_DIM = 300
+    EXPRESSION_SPACE_DIM = 100
+    NECK_IDX = 0
+
+    def __init__(
+        self,
+        model_path: str,
+        data_struct=None,
+        num_expression_coeffs=10,
+        create_expression: bool = True,
+        expression: Optional[Tensor] = None,
+        create_neck_pose: bool = True,
+        neck_pose: Optional[Tensor] = None,
+        create_jaw_pose: bool = True,
+        jaw_pose: Optional[Tensor] = None,
+        create_leye_pose: bool = True,
+        leye_pose: Optional[Tensor] = None,
+        create_reye_pose=True,
+        reye_pose: Optional[Tensor] = None,
+        use_face_contour=False,
+        batch_size: int = 1,
+        gender: str = 'neutral',
+        dtype: torch.dtype = torch.float32,
+        ext='pkl',
+        **kwargs
+    ) -> None:
+        ''' FLAME model constructor
+
+            Parameters
+            ----------
+            model_path: str
+                The path to the folder or to the file where the model
+                parameters are stored
+            num_expression_coeffs: int, optional
+                Number of expression components to use
+                (default = 10).
+            create_expression: bool, optional
+                Flag for creating a member variable for the expression space
+                (default = True).
+            expression: torch.tensor, optional, Bx10
+                The default value for the expression member variable.
+                (default = None)
+            create_neck_pose: bool, optional
+                Flag for creating a member variable for the neck pose.
+                (default = False)
+            neck_pose: torch.tensor, optional, Bx3
+                The default value for the neck pose variable.
+                (default = None)
+            create_jaw_pose: bool, optional
+                Flag for creating a member variable for the jaw pose.
+                (default = False)
+            jaw_pose: torch.tensor, optional, Bx3
+                The default value for the jaw pose variable.
+                (default = None)
+            create_leye_pose: bool, optional
+                Flag for creating a member variable for the left eye pose.
+                (default = False)
+            leye_pose: torch.tensor, optional, Bx10
+                The default value for the left eye pose variable.
+                (default = None)
+            create_reye_pose: bool, optional
+                Flag for creating a member variable for the right eye pose.
+                (default = False)
+            reye_pose: torch.tensor, optional, Bx10
+                The default value for the right eye pose variable.
+                (default = None)
+            use_face_contour: bool, optional
+                Whether to compute the keypoints that form the facial contour
+            batch_size: int, optional
+                The batch size used for creating the member variables
+            gender: str, optional
+                Which gender to load
+            dtype: torch.dtype
+                The data type for the created variables
+        '''
+        model_fn = f'FLAME_{gender.upper()}.{ext}'
+        flame_path = os.path.join(model_path, model_fn)
+        assert osp.exists(flame_path), 'Path {} does not exist!'.format(
+            flame_path)
+        if ext == 'npz':
+            file_data = np.load(flame_path, allow_pickle=True)
+        elif ext == 'pkl':
+            with open(flame_path, 'rb') as smpl_file:
+                file_data = pickle.load(smpl_file, encoding='latin1')
+        else:
+            raise ValueError('Unknown extension: {}'.format(ext))
+        data_struct = Struct(**file_data)
+
+        super(FLAME, self).__init__(
+            model_path=model_path,
+            data_struct=data_struct,
+            dtype=dtype,
+            batch_size=batch_size,
+            gender=gender,
+            ext=ext,
+            **kwargs)
+
+        self.use_face_contour = use_face_contour
+
+        self.vertex_joint_selector.extra_joints_idxs = to_tensor(
+            [], dtype=torch.long)
+
+        if create_neck_pose:
+            if neck_pose is None:
+                default_neck_pose = torch.zeros([batch_size, 3], dtype=dtype)
+            else:
+                default_neck_pose = torch.tensor(neck_pose, dtype=dtype)
+            neck_pose_param = nn.Parameter(
+                default_neck_pose, requires_grad=True)
+            self.register_parameter('neck_pose', neck_pose_param)
+
+        if create_jaw_pose:
+            if jaw_pose is None:
+                default_jaw_pose = torch.zeros([batch_size, 3], dtype=dtype)
+            else:
+                default_jaw_pose = torch.tensor(jaw_pose, dtype=dtype)
+            jaw_pose_param = nn.Parameter(default_jaw_pose,
+                                          requires_grad=True)
+            self.register_parameter('jaw_pose', jaw_pose_param)
+
+        if create_leye_pose:
+            if leye_pose is None:
+                default_leye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+            else:
+                default_leye_pose = torch.tensor(leye_pose, dtype=dtype)
+            leye_pose_param = nn.Parameter(default_leye_pose,
+                                           requires_grad=True)
+            self.register_parameter('leye_pose', leye_pose_param)
+
+        if create_reye_pose:
+            if reye_pose is None:
+                default_reye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+            else:
+                default_reye_pose = torch.tensor(reye_pose, dtype=dtype)
+            reye_pose_param = nn.Parameter(default_reye_pose,
+                                           requires_grad=True)
+            self.register_parameter('reye_pose', reye_pose_param)
+
+        shapedirs = data_struct.shapedirs
+        if len(shapedirs.shape) < 3:
+            shapedirs = shapedirs[:, :, None]
+        if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM +
+                self.EXPRESSION_SPACE_DIM):
+            print(f'WARNING: You are using a {self.name()} model, with only'
+                  ' 10 shape and 10 expression coefficients.')
+            expr_start_idx = 10
+            expr_end_idx = 20
+            num_expression_coeffs = min(num_expression_coeffs, 10)
+        else:
+            expr_start_idx = self.SHAPE_SPACE_DIM
+            expr_end_idx = self.SHAPE_SPACE_DIM + num_expression_coeffs
+            num_expression_coeffs = min(
+                num_expression_coeffs, self.EXPRESSION_SPACE_DIM)
+
+        self._num_expression_coeffs = num_expression_coeffs
+
+        expr_dirs = shapedirs[:, :, expr_start_idx:expr_end_idx]
+        self.register_buffer(
+            'expr_dirs', to_tensor(to_np(expr_dirs), dtype=dtype))
+
+        if create_expression:
+            if expression is None:
+                default_expression = torch.zeros(
+                    [batch_size, self.num_expression_coeffs], dtype=dtype)
+            else:
+                default_expression = torch.tensor(expression, dtype=dtype)
+            expression_param = nn.Parameter(default_expression,
+                                            requires_grad=True)
+            self.register_parameter('expression', expression_param)
+
+        # The pickle file that contains the barycentric coordinates for
+        # regressing the landmarks
+        landmark_bcoord_filename = osp.join(
+            model_path, 'flame_static_embedding.pkl')
+
+        with open(landmark_bcoord_filename, 'rb') as fp:
+            landmarks_data = pickle.load(fp, encoding='latin1')
+
+        lmk_faces_idx = landmarks_data['lmk_face_idx'].astype(np.int64)
+        self.register_buffer('lmk_faces_idx',
+                             torch.tensor(lmk_faces_idx, dtype=torch.long))
+        lmk_bary_coords = landmarks_data['lmk_b_coords']
+        self.register_buffer('lmk_bary_coords',
+                             torch.tensor(lmk_bary_coords, dtype=dtype))
+        if self.use_face_contour:
+            face_contour_path = os.path.join(
+                model_path, 'flame_dynamic_embedding.npy')
+            contour_embeddings = np.load(face_contour_path,
+                                         allow_pickle=True,
+                                         encoding='latin1')[()]
+
+            dynamic_lmk_faces_idx = np.array(
+                contour_embeddings['lmk_face_idx'], dtype=np.int64)
+            dynamic_lmk_faces_idx = torch.tensor(
+                dynamic_lmk_faces_idx,
+                dtype=torch.long)
+            self.register_buffer('dynamic_lmk_faces_idx',
+                                 dynamic_lmk_faces_idx)
+
+            dynamic_lmk_b_coords = torch.tensor(
+                contour_embeddings['lmk_b_coords'], dtype=dtype)
+            self.register_buffer(
+                'dynamic_lmk_bary_coords', dynamic_lmk_b_coords)
+
+            neck_kin_chain = find_joint_kin_chain(self.NECK_IDX, self.parents)
+            self.register_buffer(
+                'neck_kin_chain',
+                torch.tensor(neck_kin_chain, dtype=torch.long))
+
+    @property
+    def num_expression_coeffs(self):
+        return self._num_expression_coeffs
+
+    def name(self) -> str:
+        return 'FLAME'
+
+    def extra_repr(self):
+        msg = [
+            super(FLAME, self).extra_repr(),
+            f'Number of Expression Coefficients: {self.num_expression_coeffs}',
+            f'Use face contour: {self.use_face_contour}',
+        ]
+        return '\n'.join(msg)
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        neck_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        expression: Optional[Tensor] = None,
+        jaw_pose: Optional[Tensor] = None,
+        leye_pose: Optional[Tensor] = None,
+        reye_pose: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        pose2rot: bool = True,
+        **kwargs
+    ) -> FLAMEOutput:
+        '''
+        Forward pass for the SMPLX model
+
+            Parameters
+            ----------
+            global_orient: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable and use it as the global
+                rotation of the body. Useful if someone wishes to predicts this
+                with an external model. (default=None)
+            betas: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `betas` and use it
+                instead. For example, it can used if shape parameters
+                `betas` are predicted from some external model.
+                (default=None)
+            expression: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `expression` and use it
+                instead. For example, it can used if expression parameters
+                `expression` are predicted from some external model.
+            jaw_pose: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `jaw_pose` and
+                use this instead. It should either joint rotations in
+                axis-angle format.
+            jaw_pose: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `jaw_pose` and
+                use this instead. It should either joint rotations in
+                axis-angle format.
+            transl: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `transl` and use it
+                instead. For example, it can used if the translation
+                `transl` is predicted from some external model.
+                (default=None)
+            return_verts: bool, optional
+                Return the vertices. (default=True)
+            return_full_pose: bool, optional
+                Returns the full axis-angle pose vector (default=False)
+
+            Returns
+            -------
+                output: ModelOutput
+                A named tuple of type `ModelOutput`
+        '''
+
+        # If no shape and pose parameters are passed along, then use the
+        # ones from the module
+        global_orient = (global_orient if global_orient is not None else
+                         self.global_orient)
+        jaw_pose = jaw_pose if jaw_pose is not None else self.jaw_pose
+        neck_pose = neck_pose if neck_pose is not None else self.neck_pose
+
+        leye_pose = leye_pose if leye_pose is not None else self.leye_pose
+        reye_pose = reye_pose if reye_pose is not None else self.reye_pose
+
+        betas = betas if betas is not None else self.betas
+        expression = expression if expression is not None else self.expression
+
+        apply_trans = transl is not None or hasattr(self, 'transl')
+        if transl is None:
+            if hasattr(self, 'transl'):
+                transl = self.transl
+
+        full_pose = torch.cat(
+            [global_orient, neck_pose, jaw_pose, leye_pose, reye_pose], dim=1)
+
+        batch_size = max(betas.shape[0], global_orient.shape[0],
+                         jaw_pose.shape[0])
+        # Concatenate the shape and expression coefficients
+        scale = int(batch_size / betas.shape[0])
+        if scale > 1:
+            betas = betas.expand(scale, -1)
+        shape_components = torch.cat([betas, expression], dim=-1)
+        shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+        vertices, joints = lbs(shape_components, full_pose, self.v_template,
+                               shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=pose2rot,
+                               )
+
+        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(
+            dim=0).expand(batch_size, -1).contiguous()
+        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(
+            self.batch_size, 1, 1)
+        if self.use_face_contour:
+            lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+                vertices, full_pose, self.dynamic_lmk_faces_idx,
+                self.dynamic_lmk_bary_coords,
+                self.neck_kin_chain,
+                pose2rot=True,
+            )
+            dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+            lmk_faces_idx = torch.cat([lmk_faces_idx,
+                                       dyn_lmk_faces_idx], 1)
+            lmk_bary_coords = torch.cat(
+                [lmk_bary_coords.expand(batch_size, -1, -1),
+                 dyn_lmk_bary_coords], 1)
+
+        landmarks = vertices2landmarks(vertices, self.faces_tensor,
+                                       lmk_faces_idx,
+                                       lmk_bary_coords)
+
+        # Add any extra joints that might be needed
+        joints = self.vertex_joint_selector(vertices, joints)
+        # Add the landmarks to the joints
+        joints = torch.cat([joints, landmarks], dim=1)
+
+        # Map the joints to the current dataset
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+        if apply_trans:
+            joints += transl.unsqueeze(dim=1)
+            vertices += transl.unsqueeze(dim=1)
+
+        output = FLAMEOutput(vertices=vertices if return_verts else None,
+                             joints=joints,
+                             betas=betas,
+                             expression=expression,
+                             global_orient=global_orient,
+                             neck_pose=neck_pose,
+                             jaw_pose=jaw_pose,
+                             full_pose=full_pose if return_full_pose else None)
+        return output
+
+
+class FLAMELayer(FLAME):
+    def __init__(self, *args, **kwargs) -> None:
+        ''' FLAME as a layer model constructor '''
+        super(FLAMELayer, self).__init__(
+            create_betas=False,
+            create_expression=False,
+            create_global_orient=False,
+            create_neck_pose=False,
+            create_jaw_pose=False,
+            create_leye_pose=False,
+            create_reye_pose=False,
+            *args,
+            **kwargs)
+
+    def forward(
+        self,
+        betas: Optional[Tensor] = None,
+        global_orient: Optional[Tensor] = None,
+        neck_pose: Optional[Tensor] = None,
+        transl: Optional[Tensor] = None,
+        expression: Optional[Tensor] = None,
+        jaw_pose: Optional[Tensor] = None,
+        leye_pose: Optional[Tensor] = None,
+        reye_pose: Optional[Tensor] = None,
+        return_verts: bool = True,
+        return_full_pose: bool = False,
+        pose2rot: bool = True,
+        **kwargs
+    ) -> FLAMEOutput:
+        '''
+        Forward pass for the SMPLX model
+
+            Parameters
+            ----------
+            global_orient: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable and use it as the global
+                rotation of the body. Useful if someone wishes to predicts this
+                with an external model. (default=None)
+            betas: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `betas` and use it
+                instead. For example, it can used if shape parameters
+                `betas` are predicted from some external model.
+                (default=None)
+            expression: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `expression` and use it
+                instead. For example, it can used if expression parameters
+                `expression` are predicted from some external model.
+            jaw_pose: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `jaw_pose` and
+                use this instead. It should either joint rotations in
+                axis-angle format.
+            jaw_pose: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `jaw_pose` and
+                use this instead. It should either joint rotations in
+                axis-angle format.
+            transl: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `transl` and use it
+                instead. For example, it can used if the translation
+                `transl` is predicted from some external model.
+                (default=None)
+            return_verts: bool, optional
+                Return the vertices. (default=True)
+            return_full_pose: bool, optional
+                Returns the full axis-angle pose vector (default=False)
+
+            Returns
+            -------
+                output: ModelOutput
+                A named tuple of type `ModelOutput`
+        '''
+        device, dtype = self.shapedirs.device, self.shapedirs.dtype
+        if global_orient is None:
+            batch_size = 1
+            global_orient = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        else:
+            batch_size = global_orient.shape[0]
+        if neck_pose is None:
+            neck_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, 1, -1).contiguous()
+        if jaw_pose is None:
+            jaw_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        if leye_pose is None:
+            leye_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        if reye_pose is None:
+            reye_pose = torch.zeros(3, device=device, dtype=dtype).view(
+                1, 1, 3).expand(batch_size, -1, -1).contiguous()
+        if betas is None:
+            betas = torch.zeros([batch_size, self.num_betas],
+                                dtype=dtype, device=device)
+        if expression is None:
+            expression = torch.zeros([batch_size, self.num_expression_coeffs],
+                                     dtype=dtype, device=device)
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+        full_pose = torch.cat(
+            [global_orient, neck_pose, jaw_pose, leye_pose, reye_pose], dim=1)
+
+        shape_components = torch.cat([betas, expression], dim=-1)
+        shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+        vertices, joints = lbs(shape_components, full_pose, self.v_template,
+                               shapedirs, self.posedirs,
+                               self.J_regressor, self.parents,
+                               self.lbs_weights, pose2rot=True,
+                               )
+
+        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(
+            dim=0).expand(batch_size, -1).contiguous()
+        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(
+            self.batch_size, 1, 1)
+        if self.use_face_contour:
+            lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+                vertices, full_pose, self.dynamic_lmk_faces_idx,
+                self.dynamic_lmk_bary_coords,
+                self.neck_kin_chain,
+                pose2rot=False,
+            )
+            dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+            lmk_faces_idx = torch.cat([lmk_faces_idx,
+                                       dyn_lmk_faces_idx], 1)
+            lmk_bary_coords = torch.cat(
+                [lmk_bary_coords.expand(batch_size, -1, -1),
+                 dyn_lmk_bary_coords], 1)
+
+        landmarks = vertices2landmarks(vertices, self.faces_tensor,
+                                       lmk_faces_idx,
+                                       lmk_bary_coords)
+
+        # Add any extra joints that might be needed
+        joints = self.vertex_joint_selector(vertices, joints)
+        # Add the landmarks to the joints
+        joints = torch.cat([joints, landmarks], dim=1)
+
+        # Map the joints to the current dataset
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+        joints += transl.unsqueeze(dim=1)
+        vertices += transl.unsqueeze(dim=1)
+
+        output = FLAMEOutput(vertices=vertices if return_verts else None,
+                             joints=joints,
+                             betas=betas,
+                             expression=expression,
+                             global_orient=global_orient,
+                             neck_pose=neck_pose,
+                             jaw_pose=jaw_pose,
+                             full_pose=full_pose if return_full_pose else None)
+        return output
+
+
+def build_layer(
+    model_path: str,
+    model_type: str = 'smpl',
+    **kwargs
+) -> Union[SMPLLayer, SMPLHLayer, SMPLXLayer, MANOLayer, FLAMELayer]:
+    ''' Method for creating a model from a path and a model type
+
+        Parameters
+        ----------
+        model_path: str
+            Either the path to the model you wish to load or a folder,
+            where each subfolder contains the differents types, i.e.:
+            model_path:
+            |
+            |-- smpl
+                |-- SMPL_FEMALE
+                |-- SMPL_NEUTRAL
+                |-- SMPL_MALE
+            |-- smplh
+                |-- SMPLH_FEMALE
+                |-- SMPLH_MALE
+            |-- smplx
+                |-- SMPLX_FEMALE
+                |-- SMPLX_NEUTRAL
+                |-- SMPLX_MALE
+            |-- mano
+                |-- MANO RIGHT
+                |-- MANO LEFT
+            |-- flame
+                |-- FLAME_FEMALE
+                |-- FLAME_MALE
+                |-- FLAME_NEUTRAL
+
+        model_type: str, optional
+            When model_path is a folder, then this parameter specifies  the
+            type of model to be loaded
+        **kwargs: dict
+            Keyword arguments
+
+        Returns
+        -------
+            body_model: nn.Module
+                The PyTorch module that implements the corresponding body model
+        Raises
+        ------
+            ValueError: In case the model type is not one of SMPL, SMPLH,
+            SMPLX, MANO or FLAME
+    '''
+
+    if osp.isdir(model_path):
+        model_path = os.path.join(model_path, model_type)
+    else:
+        model_type = osp.basename(model_path).split('_')[0].lower()
+
+    if model_type.lower() == 'smpl':
+        return SMPLLayer(model_path, **kwargs)
+    elif model_type.lower() == 'smplh':
+        return SMPLHLayer(model_path, **kwargs)
+    elif model_type.lower() == 'smplx':
+        return SMPLXLayer(model_path, **kwargs)
+    elif 'mano' in model_type.lower():
+        return MANOLayer(model_path, **kwargs)
+    elif 'flame' in model_type.lower():
+        return FLAMELayer(model_path, **kwargs)
+    else:
+        raise ValueError(f'Unknown model type {model_type}, exiting!')
+
+
+def create(
+    model_path: str,
+    model_type: str = 'smpl',
+    **kwargs
+) -> Union[SMPL, SMPLH, SMPLX, MANO, FLAME]:
+    ''' Method for creating a model from a path and a model type
+
+        Parameters
+        ----------
+        model_path: str
+            Either the path to the model you wish to load or a folder,
+            where each subfolder contains the differents types, i.e.:
+            model_path:
+            |
+            |-- smpl
+                |-- SMPL_FEMALE
+                |-- SMPL_NEUTRAL
+                |-- SMPL_MALE
+            |-- smplh
+                |-- SMPLH_FEMALE
+                |-- SMPLH_MALE
+            |-- smplx
+                |-- SMPLX_FEMALE
+                |-- SMPLX_NEUTRAL
+                |-- SMPLX_MALE
+            |-- mano
+                |-- MANO RIGHT
+                |-- MANO LEFT
+
+        model_type: str, optional
+            When model_path is a folder, then this parameter specifies  the
+            type of model to be loaded
+        **kwargs: dict
+            Keyword arguments
+
+        Returns
+        -------
+            body_model: nn.Module
+                The PyTorch module that implements the corresponding body model
+        Raises
+        ------
+            ValueError: In case the model type is not one of SMPL, SMPLH,
+            SMPLX, MANO or FLAME
+    '''
+
+    # If it's a folder, assume
+    if osp.isdir(model_path):
+        model_path = os.path.join(model_path, model_type)
+    else:
+        model_type = osp.basename(model_path).split('_')[0].lower()
+
+    if model_type.lower() == 'smpl':
+        return SMPL(model_path, **kwargs)
+    elif model_type.lower() == 'smplh':
+        return SMPLH(model_path, **kwargs)
+    elif model_type.lower() == 'smplx':
+        return SMPLX(model_path, **kwargs)
+    elif 'mano' in model_type.lower():
+        return MANO(model_path, **kwargs)
+    elif 'flame' in model_type.lower():
+        return FLAME(model_path, **kwargs)
+    else:
+        raise ValueError(f'Unknown model type {model_type}, exiting!')

common/utils/smplx/smplx/joint_names.py

@@ -0,0 +1,163 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+JOINT_NAMES = [
+    'pelvis',
+    'left_hip',
+    'right_hip',
+    'spine1',
+    'left_knee',
+    'right_knee',
+    'spine2',
+    'left_ankle',
+    'right_ankle',
+    'spine3',
+    'left_foot',
+    'right_foot',
+    'neck',
+    'left_collar',
+    'right_collar',
+    'head',
+    'left_shoulder',
+    'right_shoulder',
+    'left_elbow',
+    'right_elbow',
+    'left_wrist',
+    'right_wrist',
+    'jaw',
+    'left_eye_smplhf',
+    'right_eye_smplhf',
+    'left_index1',
+    'left_index2',
+    'left_index3',
+    'left_middle1',
+    'left_middle2',
+    'left_middle3',
+    'left_pinky1',
+    'left_pinky2',
+    'left_pinky3',
+    'left_ring1',
+    'left_ring2',
+    'left_ring3',
+    'left_thumb1',
+    'left_thumb2',
+    'left_thumb3',
+    'right_index1',
+    'right_index2',
+    'right_index3',
+    'right_middle1',
+    'right_middle2',
+    'right_middle3',
+    'right_pinky1',
+    'right_pinky2',
+    'right_pinky3',
+    'right_ring1',
+    'right_ring2',
+    'right_ring3',
+    'right_thumb1',
+    'right_thumb2',
+    'right_thumb3',
+    'nose',
+    'right_eye',
+    'left_eye',
+    'right_ear',
+    'left_ear',
+    'left_big_toe',
+    'left_small_toe',
+    'left_heel',
+    'right_big_toe',
+    'right_small_toe',
+    'right_heel',
+    'left_thumb',
+    'left_index',
+    'left_middle',
+    'left_ring',
+    'left_pinky',
+    'right_thumb',
+    'right_index',
+    'right_middle',
+    'right_ring',
+    'right_pinky',
+    'right_eye_brow1',
+    'right_eye_brow2',
+    'right_eye_brow3',
+    'right_eye_brow4',
+    'right_eye_brow5',
+    'left_eye_brow5',
+    'left_eye_brow4',
+    'left_eye_brow3',
+    'left_eye_brow2',
+    'left_eye_brow1',
+    'nose1',
+    'nose2',
+    'nose3',
+    'nose4',
+    'right_nose_2',
+    'right_nose_1',
+    'nose_middle',
+    'left_nose_1',
+    'left_nose_2',
+    'right_eye1',
+    'right_eye2',
+    'right_eye3',
+    'right_eye4',
+    'right_eye5',
+    'right_eye6',
+    'left_eye4',
+    'left_eye3',
+    'left_eye2',
+    'left_eye1',
+    'left_eye6',
+    'left_eye5',
+    'right_mouth_1',
+    'right_mouth_2',
+    'right_mouth_3',
+    'mouth_top',
+    'left_mouth_3',
+    'left_mouth_2',
+    'left_mouth_1',
+    'left_mouth_5',  # 59 in OpenPose output
+    'left_mouth_4',  # 58 in OpenPose output
+    'mouth_bottom',
+    'right_mouth_4',
+    'right_mouth_5',
+    'right_lip_1',
+    'right_lip_2',
+    'lip_top',
+    'left_lip_2',
+    'left_lip_1',
+    'left_lip_3',
+    'lip_bottom',
+    'right_lip_3',
+    # Face contour
+    'right_contour_1',
+    'right_contour_2',
+    'right_contour_3',
+    'right_contour_4',
+    'right_contour_5',
+    'right_contour_6',
+    'right_contour_7',
+    'right_contour_8',
+    'contour_middle',
+    'left_contour_8',
+    'left_contour_7',
+    'left_contour_6',
+    'left_contour_5',
+    'left_contour_4',
+    'left_contour_3',
+    'left_contour_2',
+    'left_contour_1',
+]

common/utils/smplx/smplx/lbs.py

@@ -0,0 +1,404 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+from typing import Tuple, List
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+from .utils import rot_mat_to_euler, Tensor
+
+
+def find_dynamic_lmk_idx_and_bcoords(
+    vertices: Tensor,
+    pose: Tensor,
+    dynamic_lmk_faces_idx: Tensor,
+    dynamic_lmk_b_coords: Tensor,
+    neck_kin_chain: List[int],
+    pose2rot: bool = True,
+) -> Tuple[Tensor, Tensor]:
+    ''' Compute the faces, barycentric coordinates for the dynamic landmarks
+
+
+        To do so, we first compute the rotation of the neck around the y-axis
+        and then use a pre-computed look-up table to find the faces and the
+        barycentric coordinates that will be used.
+
+        Special thanks to Soubhik Sanyal ([email protected])
+        for providing the original TensorFlow implementation and for the LUT.
+
+        Parameters
+        ----------
+        vertices: torch.tensor BxVx3, dtype = torch.float32
+            The tensor of input vertices
+        pose: torch.tensor Bx(Jx3), dtype = torch.float32
+            The current pose of the body model
+        dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long
+            The look-up table from neck rotation to faces
+        dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32
+            The look-up table from neck rotation to barycentric coordinates
+        neck_kin_chain: list
+            A python list that contains the indices of the joints that form the
+            kinematic chain of the neck.
+        dtype: torch.dtype, optional
+
+        Returns
+        -------
+        dyn_lmk_faces_idx: torch.tensor, dtype = torch.long
+            A tensor of size BxL that contains the indices of the faces that
+            will be used to compute the current dynamic landmarks.
+        dyn_lmk_b_coords: torch.tensor, dtype = torch.float32
+            A tensor of size BxL that contains the indices of the faces that
+            will be used to compute the current dynamic landmarks.
+    '''
+
+    dtype = vertices.dtype
+    batch_size = vertices.shape[0]
+
+    if pose2rot:
+        aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1,
+                                     neck_kin_chain)
+        rot_mats = batch_rodrigues(
+            aa_pose.view(-1, 3)).view(batch_size, -1, 3, 3)
+    else:
+        rot_mats = torch.index_select(
+            pose.view(batch_size, -1, 3, 3), 1, neck_kin_chain)
+
+    rel_rot_mat = torch.eye(
+        3, device=vertices.device, dtype=dtype).unsqueeze_(dim=0).repeat(
+            batch_size, 1, 1)
+    for idx in range(len(neck_kin_chain)):
+        rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)
+
+    y_rot_angle = torch.round(
+        torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
+                    max=39)).to(dtype=torch.long)
+    neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
+    mask = y_rot_angle.lt(-39).to(dtype=torch.long)
+    neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
+    y_rot_angle = (neg_mask * neg_vals +
+                   (1 - neg_mask) * y_rot_angle)
+
+    dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx,
+                                           0, y_rot_angle)
+    dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords,
+                                          0, y_rot_angle)
+
+    return dyn_lmk_faces_idx, dyn_lmk_b_coords
+
+
+def vertices2landmarks(
+    vertices: Tensor,
+    faces: Tensor,
+    lmk_faces_idx: Tensor,
+    lmk_bary_coords: Tensor
+) -> Tensor:
+    ''' Calculates landmarks by barycentric interpolation
+
+        Parameters
+        ----------
+        vertices: torch.tensor BxVx3, dtype = torch.float32
+            The tensor of input vertices
+        faces: torch.tensor Fx3, dtype = torch.long
+            The faces of the mesh
+        lmk_faces_idx: torch.tensor L, dtype = torch.long
+            The tensor with the indices of the faces used to calculate the
+            landmarks.
+        lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32
+            The tensor of barycentric coordinates that are used to interpolate
+            the landmarks
+
+        Returns
+        -------
+        landmarks: torch.tensor BxLx3, dtype = torch.float32
+            The coordinates of the landmarks for each mesh in the batch
+    '''
+    # Extract the indices of the vertices for each face
+    # BxLx3
+    batch_size, num_verts = vertices.shape[:2]
+    device = vertices.device
+
+    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view(
+        batch_size, -1, 3)
+
+    lmk_faces += torch.arange(
+        batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts
+
+    lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(
+        batch_size, -1, 3, 3)
+
+    landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])
+    return landmarks
+
+
+def lbs(
+    betas: Tensor,
+    pose: Tensor,
+    v_template: Tensor,
+    shapedirs: Tensor,
+    posedirs: Tensor,
+    J_regressor: Tensor,
+    parents: Tensor,
+    lbs_weights: Tensor,
+    pose2rot: bool = True,
+) -> Tuple[Tensor, Tensor]:
+    ''' Performs Linear Blend Skinning with the given shape and pose parameters
+
+        Parameters
+        ----------
+        betas : torch.tensor BxNB
+            The tensor of shape parameters
+        pose : torch.tensor Bx(J + 1) * 3
+            The pose parameters in axis-angle format
+        v_template torch.tensor BxVx3
+            The template mesh that will be deformed
+        shapedirs : torch.tensor 1xNB
+            The tensor of PCA shape displacements
+        posedirs : torch.tensor Px(V * 3)
+            The pose PCA coefficients
+        J_regressor : torch.tensor JxV
+            The regressor array that is used to calculate the joints from
+            the position of the vertices
+        parents: torch.tensor J
+            The array that describes the kinematic tree for the model
+        lbs_weights: torch.tensor N x V x (J + 1)
+            The linear blend skinning weights that represent how much the
+            rotation matrix of each part affects each vertex
+        pose2rot: bool, optional
+            Flag on whether to convert the input pose tensor to rotation
+            matrices. The default value is True. If False, then the pose tensor
+            should already contain rotation matrices and have a size of
+            Bx(J + 1)x9
+        dtype: torch.dtype, optional
+
+        Returns
+        -------
+        verts: torch.tensor BxVx3
+            The vertices of the mesh after applying the shape and pose
+            displacements.
+        joints: torch.tensor BxJx3
+            The joints of the model
+    '''
+
+    batch_size = max(betas.shape[0], pose.shape[0])
+    device, dtype = betas.device, betas.dtype
+
+    # Add shape contribution
+    v_shaped = v_template + blend_shapes(betas, shapedirs)
+
+    # Get the joints
+    # NxJx3 array
+    J = vertices2joints(J_regressor, v_shaped)
+
+    # 3. Add pose blend shapes
+    # N x J x 3 x 3
+    ident = torch.eye(3, dtype=dtype, device=device)
+    if pose2rot:
+        rot_mats = batch_rodrigues(pose.view(-1, 3)).view(
+            [batch_size, -1, 3, 3])
+
+        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
+        # (N x P) x (P, V * 3) -> N x V x 3
+        pose_offsets = torch.matmul(
+            pose_feature, posedirs).view(batch_size, -1, 3)
+    else:
+        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
+        rot_mats = pose.view(batch_size, -1, 3, 3)
+
+        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
+                                    posedirs).view(batch_size, -1, 3)
+
+    v_posed = pose_offsets + v_shaped
+    # 4. Get the global joint location
+    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)
+
+    # 5. Do skinning:
+    # W is N x V x (J + 1)
+    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
+    # (N x V x (J + 1)) x (N x (J + 1) x 16)
+    num_joints = J_regressor.shape[0]
+    T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \
+        .view(batch_size, -1, 4, 4)
+
+    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
+                               dtype=dtype, device=device)
+    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
+    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
+
+    verts = v_homo[:, :, :3, 0]
+
+    return verts, J_transformed
+
+
+def vertices2joints(J_regressor: Tensor, vertices: Tensor) -> Tensor:
+    ''' Calculates the 3D joint locations from the vertices
+
+    Parameters
+    ----------
+    J_regressor : torch.tensor JxV
+        The regressor array that is used to calculate the joints from the
+        position of the vertices
+    vertices : torch.tensor BxVx3
+        The tensor of mesh vertices
+
+    Returns
+    -------
+    torch.tensor BxJx3
+        The location of the joints
+    '''
+
+    return torch.einsum('bik,ji->bjk', [vertices, J_regressor])
+
+
+def blend_shapes(betas: Tensor, shape_disps: Tensor) -> Tensor:
+    ''' Calculates the per vertex displacement due to the blend shapes
+
+
+    Parameters
+    ----------
+    betas : torch.tensor Bx(num_betas)
+        Blend shape coefficients
+    shape_disps: torch.tensor Vx3x(num_betas)
+        Blend shapes
+
+    Returns
+    -------
+    torch.tensor BxVx3
+        The per-vertex displacement due to shape deformation
+    '''
+
+    # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
+    # i.e. Multiply each shape displacement by its corresponding beta and
+    # then sum them.
+    blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
+    return blend_shape
+
+
+def batch_rodrigues(
+    rot_vecs: Tensor,
+    epsilon: float = 1e-8,
+) -> Tensor:
+    ''' Calculates the rotation matrices for a batch of rotation vectors
+        Parameters
+        ----------
+        rot_vecs: torch.tensor Nx3
+            array of N axis-angle vectors
+        Returns
+        -------
+        R: torch.tensor Nx3x3
+            The rotation matrices for the given axis-angle parameters
+    '''
+
+    batch_size = rot_vecs.shape[0]
+    device, dtype = rot_vecs.device, rot_vecs.dtype
+
+    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
+    rot_dir = rot_vecs / angle
+
+    cos = torch.unsqueeze(torch.cos(angle), dim=1)
+    sin = torch.unsqueeze(torch.sin(angle), dim=1)
+
+    # Bx1 arrays
+    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
+    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
+
+    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
+    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
+        .view((batch_size, 3, 3))
+
+    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
+    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
+    return rot_mat
+
+
+def transform_mat(R: Tensor, t: Tensor) -> Tensor:
+    ''' Creates a batch of transformation matrices
+        Args:
+            - R: Bx3x3 array of a batch of rotation matrices
+            - t: Bx3x1 array of a batch of translation vectors
+        Returns:
+            - T: Bx4x4 Transformation matrix
+    '''
+    # No padding left or right, only add an extra row
+    return torch.cat([F.pad(R, [0, 0, 0, 1]),
+                      F.pad(t, [0, 0, 0, 1], value=1)], dim=2)
+
+
+def batch_rigid_transform(
+    rot_mats: Tensor,
+    joints: Tensor,
+    parents: Tensor,
+    dtype=torch.float32
+) -> Tensor:
+    """
+    Applies a batch of rigid transformations to the joints
+
+    Parameters
+    ----------
+    rot_mats : torch.tensor BxNx3x3
+        Tensor of rotation matrices
+    joints : torch.tensor BxNx3
+        Locations of joints
+    parents : torch.tensor BxN
+        The kinematic tree of each object
+    dtype : torch.dtype, optional:
+        The data type of the created tensors, the default is torch.float32
+
+    Returns
+    -------
+    posed_joints : torch.tensor BxNx3
+        The locations of the joints after applying the pose rotations
+    rel_transforms : torch.tensor BxNx4x4
+        The relative (with respect to the root joint) rigid transformations
+        for all the joints
+    """
+
+    joints = torch.unsqueeze(joints, dim=-1)
+
+    rel_joints = joints.clone()
+    rel_joints[:, 1:] -= joints[:, parents[1:]]
+
+    transforms_mat = transform_mat(
+        rot_mats.reshape(-1, 3, 3),
+        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)
+
+    transform_chain = [transforms_mat[:, 0]]
+    for i in range(1, parents.shape[0]):
+        # Subtract the joint location at the rest pose
+        # No need for rotation, since it's identity when at rest
+        curr_res = torch.matmul(transform_chain[parents[i]],
+                                transforms_mat[:, i])
+        transform_chain.append(curr_res)
+
+    transforms = torch.stack(transform_chain, dim=1)
+
+    # The last column of the transformations contains the posed joints
+    posed_joints = transforms[:, :, :3, 3]
+
+    # The last column of the transformations contains the posed joints
+    posed_joints = transforms[:, :, :3, 3]
+
+    joints_homogen = F.pad(joints, [0, 0, 0, 1])
+
+    rel_transforms = transforms - F.pad(
+        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])
+
+    return posed_joints, rel_transforms

common/utils/smplx/smplx/utils.py

@@ -0,0 +1,125 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from typing import NewType, Union, Optional
+from dataclasses import dataclass, asdict, fields
+import numpy as np
+import torch
+
+Tensor = NewType('Tensor', torch.Tensor)
+Array = NewType('Array', np.ndarray)
+
+
+@dataclass
+class ModelOutput:
+    vertices: Optional[Tensor] = None
+    joints: Optional[Tensor] = None
+    full_pose: Optional[Tensor] = None
+    global_orient: Optional[Tensor] = None
+    transl: Optional[Tensor] = None
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
+    def __iter__(self):
+        return self.keys()
+
+    def keys(self):
+        keys = [t.name for t in fields(self)]
+        return iter(keys)
+
+    def values(self):
+        values = [getattr(self, t.name) for t in fields(self)]
+        return iter(values)
+
+    def items(self):
+        data = [(t.name, getattr(self, t.name)) for t in fields(self)]
+        return iter(data)
+
+
+@dataclass
+class SMPLOutput(ModelOutput):
+    betas: Optional[Tensor] = None
+    body_pose: Optional[Tensor] = None
+
+
+@dataclass
+class SMPLHOutput(SMPLOutput):
+    left_hand_pose: Optional[Tensor] = None
+    right_hand_pose: Optional[Tensor] = None
+    transl: Optional[Tensor] = None
+
+
+@dataclass
+class SMPLXOutput(SMPLHOutput):
+    expression: Optional[Tensor] = None
+    jaw_pose: Optional[Tensor] = None
+
+
+@dataclass
+class MANOOutput(ModelOutput):
+    betas: Optional[Tensor] = None
+    hand_pose: Optional[Tensor] = None
+
+
+@dataclass
+class FLAMEOutput(ModelOutput):
+    betas: Optional[Tensor] = None
+    expression: Optional[Tensor] = None
+    jaw_pose: Optional[Tensor] = None
+    neck_pose: Optional[Tensor] = None
+
+
+def find_joint_kin_chain(joint_id, kinematic_tree):
+    kin_chain = []
+    curr_idx = joint_id
+    while curr_idx != -1:
+        kin_chain.append(curr_idx)
+        curr_idx = kinematic_tree[curr_idx]
+    return kin_chain
+
+
+def to_tensor(
+        array: Union[Array, Tensor], dtype=torch.float32
+) -> Tensor:
+    if torch.is_tensor(array):
+        return array
+    else:
+        return torch.tensor(array, dtype=dtype)
+
+
+class Struct(object):
+    def __init__(self, **kwargs):
+        for key, val in kwargs.items():
+            setattr(self, key, val)
+
+
+def to_np(array, dtype=np.float32):
+    if 'scipy.sparse' in str(type(array)):
+        array = array.todense()
+    return np.array(array, dtype=dtype)
+
+
+def rot_mat_to_euler(rot_mats):
+    # Calculates rotation matrix to euler angles
+    # Careful for extreme cases of eular angles like [0.0, pi, 0.0]
+
+    sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
+                    rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
+    return torch.atan2(-rot_mats[:, 2, 0], sy)

common/utils/smplx/smplx/vertex_ids.py

@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from __future__ import print_function
+from __future__ import absolute_import
+from __future__ import division
+
+# Joint name to vertex mapping. SMPL/SMPL-H/SMPL-X vertices that correspond to
+# MSCOCO and OpenPose joints
+vertex_ids = {
+    'smplh': {
+        'nose':		    332,
+        'reye':		    6260,
+        'leye':		    2800,
+        'rear':		    4071,
+        'lear':		    583,
+        'rthumb':		6191,
+        'rindex':		5782,
+        'rmiddle':		5905,
+        'rring':		6016,
+        'rpinky':		6133,
+        'lthumb':		2746,
+        'lindex':		2319,
+        'lmiddle':		2445,
+        'lring':		2556,
+        'lpinky':		2673,
+        'LBigToe':		3216,
+        'LSmallToe':	3226,
+        'LHeel':		3387,
+        'RBigToe':		6617,
+        'RSmallToe':    6624,
+        'RHeel':		6787
+    },
+    'smplx': {
+        'nose':		    9120,
+        'reye':		    9929,
+        'leye':		    9448,
+        'rear':		    616,
+        'lear':		    6,
+        'rthumb':		8079,
+        'rindex':		7669,
+        'rmiddle':		7794,
+        'rring':		7905,
+        'rpinky':		8022,
+        'lthumb':		5361,
+        'lindex':		4933,
+        'lmiddle':		5058,
+        'lring':		5169,
+        'lpinky':		5286,
+        'LBigToe':		5770,
+        'LSmallToe':    5780,
+        'LHeel':		8846,
+        'RBigToe':		8463,
+        'RSmallToe': 	8474,
+        'RHeel':  		8635
+    },
+    'mano': {
+            'thumb':		744,
+            'index':		320,
+            'middle':		443,
+            'ring':		    554,
+            'pinky':		671,
+        }
+}

common/utils/smplx/smplx/vertex_joint_selector.py

@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from .utils import to_tensor
+
+
+class VertexJointSelector(nn.Module):
+
+    def __init__(self, vertex_ids=None,
+                 use_hands=True,
+                 use_feet_keypoints=True, **kwargs):
+        super(VertexJointSelector, self).__init__()
+
+        extra_joints_idxs = []
+
+        face_keyp_idxs = np.array([
+            vertex_ids['nose'],
+            vertex_ids['reye'],
+            vertex_ids['leye'],
+            vertex_ids['rear'],
+            vertex_ids['lear']], dtype=np.int64)
+
+        extra_joints_idxs = np.concatenate([extra_joints_idxs,
+                                            face_keyp_idxs])
+
+        if use_feet_keypoints:
+            feet_keyp_idxs = np.array([vertex_ids['LBigToe'],
+                                       vertex_ids['LSmallToe'],
+                                       vertex_ids['LHeel'],
+                                       vertex_ids['RBigToe'],
+                                       vertex_ids['RSmallToe'],
+                                       vertex_ids['RHeel']], dtype=np.int32)
+
+            extra_joints_idxs = np.concatenate(
+                [extra_joints_idxs, feet_keyp_idxs])
+
+        if use_hands:
+            self.tip_names = ['thumb', 'index', 'middle', 'ring', 'pinky']
+
+            tips_idxs = []
+            for hand_id in ['l', 'r']:
+                for tip_name in self.tip_names:
+                    tips_idxs.append(vertex_ids[hand_id + tip_name])
+
+            extra_joints_idxs = np.concatenate(
+                [extra_joints_idxs, tips_idxs])
+
+        self.register_buffer('extra_joints_idxs',
+                             to_tensor(extra_joints_idxs, dtype=torch.long))
+
+    def forward(self, vertices, joints):
+        extra_joints = torch.index_select(vertices, 1, self.extra_joints_idxs)
+        joints = torch.cat([joints, extra_joints], dim=1)
+
+        return joints

common/utils/smplx/tools/README.md

@@ -0,0 +1,20 @@
+## Removing Chumpy objects
+
+In a Python 2 virtual environment with [Chumpy](https://github.com/mattloper/chumpy) installed run the following to remove any Chumpy objects from the model data:
+
+```bash
+python tools/clean_ch.py --input-models path-to-models/*.pkl --output-folder output-folder
+```
+
+## Merging SMPL-H and MANO parameters
+
+In order to use the given PyTorch SMPL-H module we first need to merge the SMPL-H and MANO parameters in a single file. After agreeing to the license and downloading the models, run the following command:
+
+```bash
+python tools/merge_smplh_mano.py --smplh-fn SMPLH_FOLDER/SMPLH_GENDER.pkl \
+ --mano-left-fn MANO_FOLDER/MANO_LEFT.pkl \
+ --mano-right-fn MANO_FOLDER/MANO_RIGHT.pkl \
+ --output-folder OUTPUT_FOLDER
+```
+
+where SMPLH_FOLDER is the folder with the SMPL-H files and MANO_FOLDER the one for the MANO files.

common/utils/smplx/tools/__init__.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: [email protected]
+
+import clean_ch
+import merge_smplh_mano

common/utils/smplx/tools/clean_ch.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: [email protected]
+
+from __future__ import print_function
+from __future__ import absolute_import
+from __future__ import division
+
+import argparse
+import os
+import os.path as osp
+
+import pickle
+
+from tqdm import tqdm
+import numpy as np
+
+
+def clean_fn(fn, output_folder='output'):
+    with open(fn, 'rb') as body_file:
+        body_data = pickle.load(body_file)
+
+    output_dict = {}
+    for key, data in body_data.iteritems():
+        if 'chumpy' in str(type(data)):
+            output_dict[key] = np.array(data)
+        else:
+            output_dict[key] = data
+
+    out_fn = osp.split(fn)[1]
+
+    out_path = osp.join(output_folder, out_fn)
+    with open(out_path, 'wb') as out_file:
+        pickle.dump(output_dict, out_file)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input-models', dest='input_models', nargs='+',
+                        required=True, type=str,
+                        help='The path to the model that will be processed')
+    parser.add_argument('--output-folder', dest='output_folder',
+                        required=True, type=str,
+                        help='The path to the output folder')
+
+    args = parser.parse_args()
+
+    input_models = args.input_models
+    output_folder = args.output_folder
+    if not osp.exists(output_folder):
+        print('Creating directory: {}'.format(output_folder))
+        os.makedirs(output_folder)
+
+    for input_model in input_models:
+        clean_fn(input_model, output_folder=output_folder)

common/utils/smplx/tools/merge_smplh_mano.py

@@ -0,0 +1,89 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: [email protected]
+
+from __future__ import print_function
+
+import os
+import os.path as osp
+import pickle
+
+import argparse
+
+import numpy as np
+
+
+def merge_models(smplh_fn, mano_left_fn, mano_right_fn,
+                 output_folder='output'):
+
+    with open(smplh_fn, 'rb') as body_file:
+        body_data = pickle.load(body_file)
+
+    with open(mano_left_fn, 'rb') as lhand_file:
+        lhand_data = pickle.load(lhand_file)
+
+    with open(mano_right_fn, 'rb') as rhand_file:
+        rhand_data = pickle.load(rhand_file)
+
+    out_fn = osp.split(smplh_fn)[1]
+
+    output_data = body_data.copy()
+    output_data['hands_componentsl'] = lhand_data['hands_components']
+    output_data['hands_componentsr'] = rhand_data['hands_components']
+
+    output_data['hands_coeffsl'] = lhand_data['hands_coeffs']
+    output_data['hands_coeffsr'] = rhand_data['hands_coeffs']
+
+    output_data['hands_meanl'] = lhand_data['hands_mean']
+    output_data['hands_meanr'] = rhand_data['hands_mean']
+
+    for key, data in output_data.iteritems():
+        if 'chumpy' in str(type(data)):
+            output_data[key] = np.array(data)
+        else:
+            output_data[key] = data
+
+    out_path = osp.join(output_folder, out_fn)
+    print(out_path)
+    print('Saving to {}'.format(out_path))
+    with open(out_path, 'wb') as output_file:
+        pickle.dump(output_data, output_file)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--smplh-fn', dest='smplh_fn', required=True,
+                        type=str, help='The path to the SMPLH model')
+    parser.add_argument('--mano-left-fn', dest='mano_left_fn', required=True,
+                        type=str, help='The path to the left hand MANO model')
+    parser.add_argument('--mano-right-fn', dest='mano_right_fn', required=True,
+                        type=str, help='The path to the right hand MANO model')
+    parser.add_argument('--output-folder', dest='output_folder',
+                        required=True, type=str,
+                        help='The path to the output folder')
+
+    args = parser.parse_args()
+
+    smplh_fn = args.smplh_fn
+    mano_left_fn = args.mano_left_fn
+    mano_right_fn = args.mano_right_fn
+    output_folder = args.output_folder
+
+    if not osp.exists(output_folder):
+        print('Creating directory: {}'.format(output_folder))
+        os.makedirs(output_folder)
+
+    merge_models(smplh_fn, mano_left_fn, mano_right_fn, output_folder)

common/utils/transforms.py

@@ -0,0 +1,172 @@
+import torch
+import numpy as np
+import scipy
+from config import cfg
+from torch.nn import functional as F
+import torchgeometry as tgm
+
+
+def cam2pixel(cam_coord, f, c):
+    x = cam_coord[:, 0] / cam_coord[:, 2] * f[0] + c[0]
+    y = cam_coord[:, 1] / cam_coord[:, 2] * f[1] + c[1]
+    z = cam_coord[:, 2]
+    return np.stack((x, y, z), 1)
+
+
+def pixel2cam(pixel_coord, f, c):
+    x = (pixel_coord[:, 0] - c[0]) / f[0] * pixel_coord[:, 2]
+    y = (pixel_coord[:, 1] - c[1]) / f[1] * pixel_coord[:, 2]
+    z = pixel_coord[:, 2]
+    return np.stack((x, y, z), 1)
+
+
+def world2cam(world_coord, R, t):
+    cam_coord = np.dot(R, world_coord.transpose(1, 0)).transpose(1, 0) + t.reshape(1, 3)
+    return cam_coord
+
+
+def cam2world(cam_coord, R, t):
+    world_coord = np.dot(np.linalg.inv(R), (cam_coord - t.reshape(1, 3)).transpose(1, 0)).transpose(1, 0)
+    return world_coord
+
+
+def rigid_transform_3D(A, B):
+    n, dim = A.shape
+    centroid_A = np.mean(A, axis=0)
+    centroid_B = np.mean(B, axis=0)
+    H = np.dot(np.transpose(A - centroid_A), B - centroid_B) / n
+    U, s, V = np.linalg.svd(H)
+    R = np.dot(np.transpose(V), np.transpose(U))
+    if np.linalg.det(R) < 0:
+        s[-1] = -s[-1]
+        V[2] = -V[2]
+        R = np.dot(np.transpose(V), np.transpose(U))
+
+    varP = np.var(A, axis=0).sum()
+    c = 1 / varP * np.sum(s)
+
+    t = -np.dot(c * R, np.transpose(centroid_A)) + np.transpose(centroid_B)
+    return c, R, t
+
+
+def rigid_align(A, B):
+    c, R, t = rigid_transform_3D(A, B)
+    A2 = np.transpose(np.dot(c * R, np.transpose(A))) + t
+    return A2
+
+
+def transform_joint_to_other_db(src_joint, src_name, dst_name):
+    src_joint_num = len(src_name)
+    dst_joint_num = len(dst_name)
+
+    new_joint = np.zeros(((dst_joint_num,) + src_joint.shape[1:]), dtype=np.float32)
+    for src_idx in range(len(src_name)):
+        name = src_name[src_idx]
+        if name in dst_name:
+            dst_idx = dst_name.index(name)
+            new_joint[dst_idx] = src_joint[src_idx]
+
+    return new_joint
+
+
+def rot6d_to_axis_angle(x):
+    batch_size = x.shape[0]
+
+    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)
+    rot_mat = torch.stack((b1, b2, b3), dim=-1)  # 3x3 rotation matrix
+
+    rot_mat = torch.cat([rot_mat, torch.zeros((batch_size, 3, 1)).to(cfg.device).float()], 2)  # 3x4 rotation matrix
+    axis_angle = tgm.rotation_matrix_to_angle_axis(rot_mat).reshape(-1, 3)  # axis-angle
+    axis_angle[torch.isnan(axis_angle)] = 0.0
+    return axis_angle
+
+
+def sample_joint_features(img_feat, joint_xy):
+    height, width = img_feat.shape[2:]
+    x = joint_xy[:, :, 0] / (width - 1) * 2 - 1
+    y = joint_xy[:, :, 1] / (height - 1) * 2 - 1
+    grid = torch.stack((x, y), 2)[:, :, None, :]
+    img_feat = F.grid_sample(img_feat, grid, align_corners=True)[:, :, :, 0]  # batch_size, channel_dim, joint_num
+    img_feat = img_feat.permute(0, 2, 1).contiguous()  # batch_size, joint_num, channel_dim
+    return img_feat
+
+
+def soft_argmax_2d(heatmap2d):
+    batch_size = heatmap2d.shape[0]
+    height, width = heatmap2d.shape[2:]
+    heatmap2d = heatmap2d.reshape((batch_size, -1, height * width))
+    heatmap2d = F.softmax(heatmap2d, 2)
+    heatmap2d = heatmap2d.reshape((batch_size, -1, height, width))
+
+    accu_x = heatmap2d.sum(dim=(2))
+    accu_y = heatmap2d.sum(dim=(3))
+
+    accu_x = accu_x * torch.arange(width).float().to(cfg.device)[None, None, :]
+    accu_y = accu_y * torch.arange(height).float().to(cfg.device)[None, None, :]
+
+    accu_x = accu_x.sum(dim=2, keepdim=True)
+    accu_y = accu_y.sum(dim=2, keepdim=True)
+
+    coord_out = torch.cat((accu_x, accu_y), dim=2)
+    return coord_out
+
+
+def soft_argmax_3d(heatmap3d):
+    batch_size = heatmap3d.shape[0]
+    depth, height, width = heatmap3d.shape[2:]
+    heatmap3d = heatmap3d.reshape((batch_size, -1, depth * height * width))
+    heatmap3d = F.softmax(heatmap3d, 2)
+    heatmap3d = heatmap3d.reshape((batch_size, -1, depth, height, width))
+
+    accu_x = heatmap3d.sum(dim=(2, 3))
+    accu_y = heatmap3d.sum(dim=(2, 4))
+    accu_z = heatmap3d.sum(dim=(3, 4))
+
+    accu_x = accu_x * torch.arange(width).float().to(cfg.device)[None, None, :]
+    accu_y = accu_y * torch.arange(height).float().to(cfg.device)[None, None, :]
+    accu_z = accu_z * torch.arange(depth).float().to(cfg.device)[None, None, :]
+
+    accu_x = accu_x.sum(dim=2, keepdim=True)
+    accu_y = accu_y.sum(dim=2, keepdim=True)
+    accu_z = accu_z.sum(dim=2, keepdim=True)
+
+    coord_out = torch.cat((accu_x, accu_y, accu_z), dim=2)
+    return coord_out
+
+
+def restore_bbox(bbox_center, bbox_size, aspect_ratio, extension_ratio):
+    bbox = bbox_center.view(-1, 1, 2) + torch.cat((-bbox_size.view(-1, 1, 2) / 2., bbox_size.view(-1, 1, 2) / 2.),
+                                                  1)  # xyxy in (cfg.output_hm_shape[2], cfg.output_hm_shape[1]) space
+    bbox[:, :, 0] = bbox[:, :, 0] / cfg.output_hm_shape[2] * cfg.input_body_shape[1]
+    bbox[:, :, 1] = bbox[:, :, 1] / cfg.output_hm_shape[1] * cfg.input_body_shape[0]
+    bbox = bbox.view(-1, 4)
+
+    # xyxy -> xywh
+    bbox[:, 2] = bbox[:, 2] - bbox[:, 0]
+    bbox[:, 3] = bbox[:, 3] - bbox[:, 1]
+
+    # aspect ratio preserving bbox
+    w = bbox[:, 2]
+    h = bbox[:, 3]
+    c_x = bbox[:, 0] + w / 2.
+    c_y = bbox[:, 1] + h / 2.
+
+    mask1 = w > (aspect_ratio * h)
+    mask2 = w < (aspect_ratio * h)
+    h[mask1] = w[mask1] / aspect_ratio
+    w[mask2] = h[mask2] * aspect_ratio
+
+    bbox[:, 2] = w * extension_ratio
+    bbox[:, 3] = h * extension_ratio
+    bbox[:, 0] = c_x - bbox[:, 2] / 2.
+    bbox[:, 1] = c_y - bbox[:, 3] / 2.
+
+    # xywh -> xyxy
+    bbox[:, 2] = bbox[:, 2] + bbox[:, 0]
+    bbox[:, 3] = bbox[:, 3] + bbox[:, 1]
+    return bbox

common/utils/vis.py

@@ -0,0 +1,183 @@
+import os
+import cv2
+import numpy as np
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+import os
+os.environ["PYOPENGL_PLATFORM"] = "egl"
+import pyrender
+import trimesh
+from config import cfg
+
+def vis_keypoints_with_skeleton(img, kps, kps_lines, kp_thresh=0.4, alpha=1):
+    # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
+    cmap = plt.get_cmap('rainbow')
+    colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]
+    colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
+
+    # Perform the drawing on a copy of the image, to allow for blending.
+    kp_mask = np.copy(img)
+
+    # Draw the keypoints.
+    for l in range(len(kps_lines)):
+        i1 = kps_lines[l][0]
+        i2 = kps_lines[l][1]
+        p1 = kps[0, i1].astype(np.int32), kps[1, i1].astype(np.int32)
+        p2 = kps[0, i2].astype(np.int32), kps[1, i2].astype(np.int32)
+        if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
+            cv2.line(
+                kp_mask, p1, p2,
+                color=colors[l], thickness=2, lineType=cv2.LINE_AA)
+        if kps[2, i1] > kp_thresh:
+            cv2.circle(
+                kp_mask, p1,
+                radius=3, color=colors[l], thickness=-1, lineType=cv2.LINE_AA)
+        if kps[2, i2] > kp_thresh:
+            cv2.circle(
+                kp_mask, p2,
+                radius=3, color=colors[l], thickness=-1, lineType=cv2.LINE_AA)
+
+    # Blend the keypoints.
+    return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)
+
+def vis_keypoints(img, kps, alpha=1, radius=3, color=None):
+    # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
+    cmap = plt.get_cmap('rainbow')
+    if color is None:
+        colors = [cmap(i) for i in np.linspace(0, 1, len(kps) + 2)]
+        colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
+
+    # Perform the drawing on a copy of the image, to allow for blending.
+    kp_mask = np.copy(img)
+
+    # Draw the keypoints.
+    for i in range(len(kps)):
+        p = kps[i][0].astype(np.int32), kps[i][1].astype(np.int32)
+        if color is None:
+            cv2.circle(kp_mask, p, radius=radius, color=colors[i], thickness=-1, lineType=cv2.LINE_AA)
+        else:
+            cv2.circle(kp_mask, p, radius=radius, color=color, thickness=-1, lineType=cv2.LINE_AA)
+
+    # Blend the keypoints.
+    return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)
+
+def vis_mesh(img, mesh_vertex, alpha=0.5):
+    # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
+    cmap = plt.get_cmap('rainbow')
+    colors = [cmap(i) for i in np.linspace(0, 1, len(mesh_vertex))]
+    colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
+
+    # Perform the drawing on a copy of the image, to allow for blending.
+    mask = np.copy(img)
+
+    # Draw the mesh
+    for i in range(len(mesh_vertex)):
+        p = mesh_vertex[i][0].astype(np.int32), mesh_vertex[i][1].astype(np.int32)
+        cv2.circle(mask, p, radius=1, color=colors[i], thickness=-1, lineType=cv2.LINE_AA)
+
+    # Blend the keypoints.
+    return cv2.addWeighted(img, 1.0 - alpha, mask, alpha, 0)
+
+def vis_3d_skeleton(kpt_3d, kpt_3d_vis, kps_lines, filename=None):
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+
+    # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
+    cmap = plt.get_cmap('rainbow')
+    colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]
+    colors = [np.array((c[2], c[1], c[0])) for c in colors]
+
+    for l in range(len(kps_lines)):
+        i1 = kps_lines[l][0]
+        i2 = kps_lines[l][1]
+        x = np.array([kpt_3d[i1,0], kpt_3d[i2,0]])
+        y = np.array([kpt_3d[i1,1], kpt_3d[i2,1]])
+        z = np.array([kpt_3d[i1,2], kpt_3d[i2,2]])
+
+        if kpt_3d_vis[i1,0] > 0 and kpt_3d_vis[i2,0] > 0:
+            ax.plot(x, z, -y, c=colors[l], linewidth=2)
+        if kpt_3d_vis[i1,0] > 0:
+            ax.scatter(kpt_3d[i1,0], kpt_3d[i1,2], -kpt_3d[i1,1], c=colors[l], marker='o')
+        if kpt_3d_vis[i2,0] > 0:
+            ax.scatter(kpt_3d[i2,0], kpt_3d[i2,2], -kpt_3d[i2,1], c=colors[l], marker='o')
+
+    x_r = np.array([0, cfg.input_shape[1]], dtype=np.float32)
+    y_r = np.array([0, cfg.input_shape[0]], dtype=np.float32)
+    z_r = np.array([0, 1], dtype=np.float32)
+    
+    if filename is None:
+        ax.set_title('3D vis')
+    else:
+        ax.set_title(filename)
+
+    ax.set_xlabel('X Label')
+    ax.set_ylabel('Z Label')
+    ax.set_zlabel('Y Label')
+    ax.legend()
+
+    plt.show()
+    cv2.waitKey(0)
+
+def save_obj(v, f, file_name='output.obj'):
+    obj_file = open(file_name, 'w')
+    for i in range(len(v)):
+        obj_file.write('v ' + str(v[i][0]) + ' ' + str(v[i][1]) + ' ' + str(v[i][2]) + '\n')
+    for i in range(len(f)):
+        obj_file.write('f ' + str(f[i][0]+1) + '/' + str(f[i][0]+1) + ' ' + str(f[i][1]+1) + '/' + str(f[i][1]+1) + ' ' + str(f[i][2]+1) + '/' + str(f[i][2]+1) + '\n')
+    obj_file.close()
+
+
+def perspective_projection(vertices, cam_param):
+    # vertices: [N, 3]
+    # cam_param: [3]
+    fx, fy= cam_param['focal']
+    cx, cy = cam_param['princpt']
+    vertices[:, 0] = vertices[:, 0] * fx / vertices[:, 2] + cx
+    vertices[:, 1] = vertices[:, 1] * fy / vertices[:, 2] + cy
+    return vertices
+
+
+def render_mesh(img, mesh, face, cam_param, mesh_as_vertices=False):
+    if mesh_as_vertices:
+        # to run on cluster where headless pyrender is not supported for A100/V100
+        vertices_2d = perspective_projection(mesh, cam_param)
+        img = vis_keypoints(img, vertices_2d, alpha=0.8, radius=2, color=(0, 0, 255))
+    else:
+        # mesh
+        mesh = trimesh.Trimesh(mesh, face)
+        rot = trimesh.transformations.rotation_matrix(
+        np.radians(180), [1, 0, 0])
+        mesh.apply_transform(rot)
+        material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.0, alphaMode='OPAQUE', baseColorFactor=(1.0, 1.0, 0.9, 1.0))
+        mesh = pyrender.Mesh.from_trimesh(mesh, material=material, smooth=False)
+        scene = pyrender.Scene(ambient_light=(0.3, 0.3, 0.3))
+        scene.add(mesh, 'mesh')
+
+        focal, princpt = cam_param['focal'], cam_param['princpt']
+        camera = pyrender.IntrinsicsCamera(fx=focal[0], fy=focal[1], cx=princpt[0], cy=princpt[1])
+        scene.add(camera)
+
+        # renderer
+        renderer = pyrender.OffscreenRenderer(viewport_width=img.shape[1], viewport_height=img.shape[0], point_size=1.0)
+
+        # light
+        light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=0.8)
+        light_pose = np.eye(4)
+        light_pose[:3, 3] = np.array([0, -1, 1])
+        scene.add(light, pose=light_pose)
+        light_pose[:3, 3] = np.array([0, 1, 1])
+        scene.add(light, pose=light_pose)
+        light_pose[:3, 3] = np.array([1, 1, 2])
+        scene.add(light, pose=light_pose)
+
+        # render
+        rgb, depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
+        rgb = rgb[:,:,:3].astype(np.float32)
+        valid_mask = (depth > 0)[:,:,None]
+
+        # save to image
+        img = rgb * valid_mask + img * (1-valid_mask)
+
+    return img

main/SMPLer_X.py

@@ -0,0 +1,468 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from nets.smpler_x import PositionNet, HandRotationNet, FaceRegressor, BoxNet, HandRoI, BodyRotationNet
+from nets.loss import CoordLoss, ParamLoss, CELoss
+from utils.human_models import smpl_x
+from utils.transforms import rot6d_to_axis_angle, restore_bbox
+from config import cfg
+import math
+import copy
+from mmpose.models import build_posenet
+from mmcv import Config
+
+class Model(nn.Module):
+    def __init__(self, encoder, body_position_net, body_rotation_net, box_net, hand_position_net, hand_roi_net,
+                 hand_rotation_net, face_regressor):
+        super(Model, self).__init__()
+
+        # body
+        self.encoder = encoder
+        self.body_position_net = body_position_net
+        self.body_regressor = body_rotation_net
+        self.box_net = box_net
+
+        # hand
+        self.hand_roi_net = hand_roi_net
+        self.hand_position_net = hand_position_net
+        self.hand_regressor = hand_rotation_net
+
+        # face
+        self.face_regressor = face_regressor
+
+        self.smplx_layer = copy.deepcopy(smpl_x.layer['neutral']).to(cfg.device)
+        self.coord_loss = CoordLoss()
+        self.param_loss = ParamLoss()
+        self.ce_loss = CELoss()
+
+        self.body_num_joints = len(smpl_x.pos_joint_part['body'])
+        self.hand_joint_num = len(smpl_x.pos_joint_part['rhand'])
+
+        self.neck = [self.box_net, self.hand_roi_net]
+
+        self.head = [self.body_position_net, self.body_regressor,
+                    self.hand_position_net, self.hand_regressor, 
+                    self.face_regressor]
+
+        self.trainable_modules = [self.encoder, self.body_position_net, self.body_regressor,
+                                  self.box_net, self.hand_position_net,
+                                  self.hand_roi_net, self.hand_regressor, self.face_regressor]
+        self.special_trainable_modules = []
+
+        # backbone:
+        param_bb = sum(p.numel() for p in self.encoder.parameters() if p.requires_grad)
+        # neck 
+        param_neck = 0
+        for module in self.neck:
+            param_neck += sum(p.numel() for p in module.parameters() if p.requires_grad)
+        # head
+        param_head = 0
+        for module in self.head:
+            param_head += sum(p.numel() for p in module.parameters() if p.requires_grad)
+
+        param_net = param_bb + param_neck + param_head
+
+        # print('#parameters:')
+        # print(f'{param_bb}, {param_neck}, {param_head}, {param_net}')
+
+    def get_camera_trans(self, cam_param):
+        # camera translation
+        t_xy = cam_param[:, :2]
+        gamma = torch.sigmoid(cam_param[:, 2])  # apply sigmoid to make it positive
+        k_value = torch.FloatTensor([math.sqrt(cfg.focal[0] * cfg.focal[1] * cfg.camera_3d_size * cfg.camera_3d_size / (
+                cfg.input_body_shape[0] * cfg.input_body_shape[1]))]).to(cfg.device).view(-1)
+        t_z = k_value * gamma
+        cam_trans = torch.cat((t_xy, t_z[:, None]), 1)
+        return cam_trans
+
+    def get_coord(self, root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose, shape, expr, cam_trans, mode):
+        batch_size = root_pose.shape[0]
+        zero_pose = torch.zeros((1, 3)).float().to(cfg.device).repeat(batch_size, 1)  # eye poses
+        output = self.smplx_layer(betas=shape, body_pose=body_pose, global_orient=root_pose, right_hand_pose=rhand_pose,
+                                  left_hand_pose=lhand_pose, jaw_pose=jaw_pose, leye_pose=zero_pose,
+                                  reye_pose=zero_pose, expression=expr)
+        # camera-centered 3D coordinate
+        mesh_cam = output.vertices
+        if mode == 'test' and cfg.testset == 'AGORA':  # use 144 joints for AGORA evaluation
+            joint_cam = output.joints
+        else:
+            joint_cam = output.joints[:, smpl_x.joint_idx, :]
+
+        # project 3D coordinates to 2D space
+        if mode == 'train' and len(cfg.trainset_3d) == 1 and cfg.trainset_3d[0] == 'AGORA' and len(
+                cfg.trainset_2d) == 0:  # prevent gradients from backpropagating to SMPLX paraemter regression module
+            x = (joint_cam[:, :, 0].detach() + cam_trans[:, None, 0]) / (
+                    joint_cam[:, :, 2].detach() + cam_trans[:, None, 2] + 1e-4) * cfg.focal[0] + cfg.princpt[0]
+            y = (joint_cam[:, :, 1].detach() + cam_trans[:, None, 1]) / (
+                    joint_cam[:, :, 2].detach() + cam_trans[:, None, 2] + 1e-4) * cfg.focal[1] + cfg.princpt[1]
+        else:
+            x = (joint_cam[:, :, 0] + cam_trans[:, None, 0]) / (joint_cam[:, :, 2] + cam_trans[:, None, 2] + 1e-4) * \
+                cfg.focal[0] + cfg.princpt[0]
+            y = (joint_cam[:, :, 1] + cam_trans[:, None, 1]) / (joint_cam[:, :, 2] + cam_trans[:, None, 2] + 1e-4) * \
+                cfg.focal[1] + cfg.princpt[1]
+        x = x / cfg.input_body_shape[1] * cfg.output_hm_shape[2]
+        y = y / cfg.input_body_shape[0] * cfg.output_hm_shape[1]
+        joint_proj = torch.stack((x, y), 2)
+
+        # root-relative 3D coordinates
+        root_cam = joint_cam[:, smpl_x.root_joint_idx, None, :]
+        joint_cam = joint_cam - root_cam
+        mesh_cam = mesh_cam + cam_trans[:, None, :]  # for rendering
+        joint_cam_wo_ra = joint_cam.clone()
+
+        # left hand root (left wrist)-relative 3D coordinatese
+        lhand_idx = smpl_x.joint_part['lhand']
+        lhand_cam = joint_cam[:, lhand_idx, :]
+        lwrist_cam = joint_cam[:, smpl_x.lwrist_idx, None, :]
+        lhand_cam = lhand_cam - lwrist_cam
+        joint_cam = torch.cat((joint_cam[:, :lhand_idx[0], :], lhand_cam, joint_cam[:, lhand_idx[-1] + 1:, :]), 1)
+
+        # right hand root (right wrist)-relative 3D coordinatese
+        rhand_idx = smpl_x.joint_part['rhand']
+        rhand_cam = joint_cam[:, rhand_idx, :]
+        rwrist_cam = joint_cam[:, smpl_x.rwrist_idx, None, :]
+        rhand_cam = rhand_cam - rwrist_cam
+        joint_cam = torch.cat((joint_cam[:, :rhand_idx[0], :], rhand_cam, joint_cam[:, rhand_idx[-1] + 1:, :]), 1)
+
+        # face root (neck)-relative 3D coordinates
+        face_idx = smpl_x.joint_part['face']
+        face_cam = joint_cam[:, face_idx, :]
+        neck_cam = joint_cam[:, smpl_x.neck_idx, None, :]
+        face_cam = face_cam - neck_cam
+        joint_cam = torch.cat((joint_cam[:, :face_idx[0], :], face_cam, joint_cam[:, face_idx[-1] + 1:, :]), 1)
+
+        return joint_proj, joint_cam, joint_cam_wo_ra, mesh_cam
+
+    def generate_mesh_gt(self, targets, mode):
+        if 'smplx_mesh_cam' in targets:
+            return targets['smplx_mesh_cam']
+        nums = [3, 63, 45, 45, 3]
+        accu = []
+        temp = 0
+        for num in nums:
+            temp += num
+            accu.append(temp)
+        pose = targets['smplx_pose']
+        root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose = \
+            pose[:, :accu[0]], pose[:, accu[0]:accu[1]], pose[:, accu[1]:accu[2]], pose[:, accu[2]:accu[3]], pose[:,
+                                                                                                             accu[3]:
+                                                                                                             accu[4]]
+        # print(lhand_pose)
+        shape = targets['smplx_shape']
+        expr = targets['smplx_expr']
+        cam_trans = targets['smplx_cam_trans']
+
+        # final output
+        joint_proj, joint_cam, joint_cam_wo_ra, mesh_cam = self.get_coord(root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose, shape,
+                                                         expr, cam_trans, mode)
+
+        return mesh_cam
+
+    def bbox_split(self, bbox):
+        # bbox:[bs, 3, 3]
+        lhand_bbox_center, rhand_bbox_center, face_bbox_center = \
+            bbox[:, 0, :2], bbox[:, 1, :2], bbox[:, 2, :2]
+        return lhand_bbox_center, rhand_bbox_center, face_bbox_center
+
+    def forward(self, inputs, targets, meta_info, mode):
+
+        body_img = F.interpolate(inputs['img'], cfg.input_body_shape)
+
+        # 1. Encoder
+        img_feat, task_tokens = self.encoder(body_img)  # task_token:[bs, N, c]
+        shape_token, cam_token, expr_token, jaw_pose_token, hand_token, body_pose_token = \
+            task_tokens[:, 0], task_tokens[:, 1], task_tokens[:, 2], task_tokens[:, 3], task_tokens[:, 4:6], task_tokens[:, 6:]
+
+        # 2. Body Regressor
+        body_joint_hm, body_joint_img = self.body_position_net(img_feat)
+        root_pose, body_pose, shape, cam_param, = self.body_regressor(body_pose_token, shape_token, cam_token, body_joint_img.detach())
+        root_pose = rot6d_to_axis_angle(root_pose)
+        body_pose = rot6d_to_axis_angle(body_pose.reshape(-1, 6)).reshape(body_pose.shape[0], -1)  # (N, J_R*3)
+        cam_trans = self.get_camera_trans(cam_param)
+
+        # 3. Hand and Face BBox Estimation
+        lhand_bbox_center, lhand_bbox_size, rhand_bbox_center, rhand_bbox_size, face_bbox_center, face_bbox_size = self.box_net(img_feat, body_joint_hm.detach())
+        lhand_bbox = restore_bbox(lhand_bbox_center, lhand_bbox_size, cfg.input_hand_shape[1] / cfg.input_hand_shape[0], 2.0).detach()  # xyxy in (cfg.input_body_shape[1], cfg.input_body_shape[0]) space
+        rhand_bbox = restore_bbox(rhand_bbox_center, rhand_bbox_size, cfg.input_hand_shape[1] / cfg.input_hand_shape[0], 2.0).detach()  # xyxy in (cfg.input_body_shape[1], cfg.input_body_shape[0]) space
+        face_bbox = restore_bbox(face_bbox_center, face_bbox_size, cfg.input_face_shape[1] / cfg.input_face_shape[0], 1.5).detach()  # xyxy in (cfg.input_body_shape[1], cfg.input_body_shape[0]) space
+
+        # 4. Differentiable Feature-level Hand Crop-Upsample
+        # hand_feat: list, [bsx2, c, cfg.output_hm_shape[1]*scale, cfg.output_hm_shape[2]*scale]
+        hand_feat = self.hand_roi_net(img_feat, lhand_bbox, rhand_bbox)  # hand_feat: flipped left hand + right hand
+
+        # 5. Hand/Face Regressor
+        # hand regressor
+        _, hand_joint_img = self.hand_position_net(hand_feat)  # (2N, J_P, 3)
+        hand_pose = self.hand_regressor(hand_feat, hand_joint_img.detach())
+        hand_pose = rot6d_to_axis_angle(hand_pose.reshape(-1, 6)).reshape(hand_feat.shape[0], -1)  # (2N, J_R*3)
+        # restore flipped left hand joint coordinates
+        batch_size = hand_joint_img.shape[0] // 2
+        lhand_joint_img = hand_joint_img[:batch_size, :, :]
+        lhand_joint_img = torch.cat((cfg.output_hand_hm_shape[2] - 1 - lhand_joint_img[:, :, 0:1], lhand_joint_img[:, :, 1:]), 2)
+        rhand_joint_img = hand_joint_img[batch_size:, :, :]
+        # restore flipped left hand joint rotations
+        batch_size = hand_pose.shape[0] // 2
+        lhand_pose = hand_pose[:batch_size, :].reshape(-1, len(smpl_x.orig_joint_part['lhand']), 3)
+        lhand_pose = torch.cat((lhand_pose[:, :, 0:1], -lhand_pose[:, :, 1:3]), 2).view(batch_size, -1)
+        rhand_pose = hand_pose[batch_size:, :]
+
+        # hand regressor
+        expr, jaw_pose = self.face_regressor(expr_token, jaw_pose_token)
+        jaw_pose = rot6d_to_axis_angle(jaw_pose)
+
+        # final output
+        joint_proj, joint_cam, joint_cam_wo_ra, mesh_cam = self.get_coord(root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose, shape, expr, cam_trans, mode)
+        pose = torch.cat((root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose), 1)
+        joint_img = torch.cat((body_joint_img, lhand_joint_img, rhand_joint_img), 1)
+
+        if mode == 'test' and 'smplx_pose' in targets:
+            mesh_pseudo_gt = self.generate_mesh_gt(targets, mode)
+
+        if mode == 'train':
+            # loss functions
+            loss = {}
+
+            smplx_kps_3d_weight = getattr(cfg, 'smplx_kps_3d_weight', 1.0)
+            smplx_kps_3d_weight = getattr(cfg, 'smplx_kps_weight', smplx_kps_3d_weight) # old config
+
+            smplx_kps_2d_weight = getattr(cfg, 'smplx_kps_2d_weight', 1.0)
+            net_kps_2d_weight = getattr(cfg, 'net_kps_2d_weight', 1.0)
+
+            smplx_pose_weight = getattr(cfg, 'smplx_pose_weight', 1.0)
+            smplx_shape_weight = getattr(cfg, 'smplx_loss_weight', 1.0)
+            # smplx_orient_weight = getattr(cfg, 'smplx_orient_weight', smplx_pose_weight) # if not specified, use the same weight as pose
+    
+
+            # do not supervise root pose if original agora json is used
+            if getattr(cfg, 'agora_fix_global_orient_transl', False):
+                # loss['smplx_pose'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid'])[:, 3:] * smplx_pose_weight
+                if hasattr(cfg, 'smplx_orient_weight'):
+                    smplx_orient_weight = getattr(cfg, 'smplx_orient_weight')
+                    loss['smplx_orient'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid'])[:, :3] * smplx_orient_weight
+
+                loss['smplx_pose'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid']) * smplx_pose_weight
+
+            else:
+                loss['smplx_pose'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid'])[:, 3:] * smplx_pose_weight
+
+            loss['smplx_shape'] = self.param_loss(shape, targets['smplx_shape'],
+                                                  meta_info['smplx_shape_valid'][:, None]) * smplx_shape_weight 
+            loss['smplx_expr'] = self.param_loss(expr, targets['smplx_expr'], meta_info['smplx_expr_valid'][:, None])
+
+            # supervision for keypoints3d wo/ ra
+            loss['joint_cam'] = self.coord_loss(joint_cam_wo_ra, targets['joint_cam'], meta_info['joint_valid'] * meta_info['is_3D'][:, None, None]) * smplx_kps_3d_weight
+            # supervision for keypoints3d w/ ra
+            loss['smplx_joint_cam'] = self.coord_loss(joint_cam, targets['smplx_joint_cam'], meta_info['smplx_joint_valid']) * smplx_kps_3d_weight
+
+            if not (meta_info['lhand_bbox_valid'] == 0).all():
+                loss['lhand_bbox'] = (self.coord_loss(lhand_bbox_center, targets['lhand_bbox_center'], meta_info['lhand_bbox_valid'][:, None]) +
+                                    self.coord_loss(lhand_bbox_size, targets['lhand_bbox_size'], meta_info['lhand_bbox_valid'][:, None]))
+            if not (meta_info['rhand_bbox_valid'] == 0).all():
+                loss['rhand_bbox'] = (self.coord_loss(rhand_bbox_center, targets['rhand_bbox_center'], meta_info['rhand_bbox_valid'][:, None]) +
+                                    self.coord_loss(rhand_bbox_size, targets['rhand_bbox_size'], meta_info['rhand_bbox_valid'][:, None]))
+            if not (meta_info['face_bbox_valid'] == 0).all():
+                loss['face_bbox'] = (self.coord_loss(face_bbox_center, targets['face_bbox_center'], meta_info['face_bbox_valid'][:, None]) +
+                                 self.coord_loss(face_bbox_size, targets['face_bbox_size'], meta_info['face_bbox_valid'][:, None]))
+            
+            # if (meta_info['face_bbox_valid'] == 0).all():
+            #     out = {}
+            targets['original_joint_img'] = targets['joint_img'].clone()
+            targets['original_smplx_joint_img'] = targets['smplx_joint_img'].clone()
+            # out['original_joint_proj'] = joint_proj.clone()
+            if not (meta_info['lhand_bbox_valid'] + meta_info['rhand_bbox_valid'] == 0).all():
+
+                # change hand target joint_img and joint_trunc according to hand bbox (cfg.output_hm_shape -> downsampled hand bbox space)
+                for part_name, bbox in (('lhand', lhand_bbox), ('rhand', rhand_bbox)):
+                    for coord_name, trunc_name in (('joint_img', 'joint_trunc'), ('smplx_joint_img', 'smplx_joint_trunc')):
+                        x = targets[coord_name][:, smpl_x.joint_part[part_name], 0]
+                        y = targets[coord_name][:, smpl_x.joint_part[part_name], 1]
+                        z = targets[coord_name][:, smpl_x.joint_part[part_name], 2]
+                        trunc = meta_info[trunc_name][:, smpl_x.joint_part[part_name], 0]
+
+                        x -= (bbox[:, None, 0] / cfg.input_body_shape[1] * cfg.output_hm_shape[2])
+                        x *= (cfg.output_hand_hm_shape[2] / (
+                                (bbox[:, None, 2] - bbox[:, None, 0]) / cfg.input_body_shape[1] * cfg.output_hm_shape[
+                            2]))
+                        y -= (bbox[:, None, 1] / cfg.input_body_shape[0] * cfg.output_hm_shape[1])
+                        y *= (cfg.output_hand_hm_shape[1] / (
+                                (bbox[:, None, 3] - bbox[:, None, 1]) / cfg.input_body_shape[0] * cfg.output_hm_shape[
+                            1]))
+                        z *= cfg.output_hand_hm_shape[0] / cfg.output_hm_shape[0]
+                        trunc *= ((x >= 0) * (x < cfg.output_hand_hm_shape[2]) * (y >= 0) * (
+                                y < cfg.output_hand_hm_shape[1]))
+
+                        coord = torch.stack((x, y, z), 2)
+                        trunc = trunc[:, :, None]
+                        targets[coord_name] = torch.cat((targets[coord_name][:, :smpl_x.joint_part[part_name][0], :], coord,
+                                                        targets[coord_name][:, smpl_x.joint_part[part_name][-1] + 1:, :]),
+                                                        1)
+                        meta_info[trunc_name] = torch.cat((meta_info[trunc_name][:, :smpl_x.joint_part[part_name][0], :],
+                                                        trunc,
+                                                        meta_info[trunc_name][:, smpl_x.joint_part[part_name][-1] + 1:,
+                                                        :]), 1)
+
+                # change hand projected joint coordinates according to hand bbox (cfg.output_hm_shape -> hand bbox space)
+                for part_name, bbox in (('lhand', lhand_bbox), ('rhand', rhand_bbox)):
+                    x = joint_proj[:, smpl_x.joint_part[part_name], 0]
+                    y = joint_proj[:, smpl_x.joint_part[part_name], 1]
+
+                    x -= (bbox[:, None, 0] / cfg.input_body_shape[1] * cfg.output_hm_shape[2])
+                    x *= (cfg.output_hand_hm_shape[2] / (
+                            (bbox[:, None, 2] - bbox[:, None, 0]) / cfg.input_body_shape[1] * cfg.output_hm_shape[2]))
+                    y -= (bbox[:, None, 1] / cfg.input_body_shape[0] * cfg.output_hm_shape[1])
+                    y *= (cfg.output_hand_hm_shape[1] / (
+                            (bbox[:, None, 3] - bbox[:, None, 1]) / cfg.input_body_shape[0] * cfg.output_hm_shape[1]))
+
+                    coord = torch.stack((x, y), 2)
+                    trans = []
+                    for bid in range(coord.shape[0]):
+                        mask = meta_info['joint_trunc'][bid, smpl_x.joint_part[part_name], 0] == 1
+                        if torch.sum(mask) == 0:
+                            trans.append(torch.zeros((2)).float().to(cfg.device))
+                        else:
+                            trans.append((-coord[bid, mask, :2] + targets['joint_img'][:, smpl_x.joint_part[part_name], :][
+                                                                bid, mask, :2]).mean(0))
+                    trans = torch.stack(trans)[:, None, :]
+                    coord = coord + trans  # global translation alignment
+                    joint_proj = torch.cat((joint_proj[:, :smpl_x.joint_part[part_name][0], :], coord,
+                                            joint_proj[:, smpl_x.joint_part[part_name][-1] + 1:, :]), 1)
+
+            if not (meta_info['face_bbox_valid'] == 0).all():
+                # change face projected joint coordinates according to face bbox (cfg.output_hm_shape -> face bbox space)
+                coord = joint_proj[:, smpl_x.joint_part['face'], :]
+                trans = []
+                for bid in range(coord.shape[0]):
+                    mask = meta_info['joint_trunc'][bid, smpl_x.joint_part['face'], 0] == 1
+                    if torch.sum(mask) == 0:
+                        trans.append(torch.zeros((2)).float().to(cfg.device))
+                    else:
+                        trans.append((-coord[bid, mask, :2] + targets['joint_img'][:, smpl_x.joint_part['face'], :][bid,
+                                                            mask, :2]).mean(0))
+                trans = torch.stack(trans)[:, None, :]
+                coord = coord + trans  # global translation alignment
+                joint_proj = torch.cat((joint_proj[:, :smpl_x.joint_part['face'][0], :], coord,
+                                        joint_proj[:, smpl_x.joint_part['face'][-1] + 1:, :]), 1)
+            
+            loss['joint_proj'] = self.coord_loss(joint_proj, targets['joint_img'][:, :, :2], meta_info['joint_trunc']) * smplx_kps_2d_weight
+            loss['joint_img'] = self.coord_loss(joint_img, smpl_x.reduce_joint_set(targets['joint_img']),
+                                                smpl_x.reduce_joint_set(meta_info['joint_trunc']), meta_info['is_3D']) * net_kps_2d_weight
+            
+            loss['smplx_joint_img'] = self.coord_loss(joint_img, smpl_x.reduce_joint_set(targets['smplx_joint_img']),
+                                                      smpl_x.reduce_joint_set(meta_info['smplx_joint_trunc'])) * net_kps_2d_weight
+
+            return loss
+        else:
+            # change hand output joint_img according to hand bbox
+            for part_name, bbox in (('lhand', lhand_bbox), ('rhand', rhand_bbox)):
+                joint_img[:, smpl_x.pos_joint_part[part_name], 0] *= (
+                        ((bbox[:, None, 2] - bbox[:, None, 0]) / cfg.input_body_shape[1] * cfg.output_hm_shape[2]) /
+                        cfg.output_hand_hm_shape[2])
+                joint_img[:, smpl_x.pos_joint_part[part_name], 0] += (
+                        bbox[:, None, 0] / cfg.input_body_shape[1] * cfg.output_hm_shape[2])
+                joint_img[:, smpl_x.pos_joint_part[part_name], 1] *= (
+                        ((bbox[:, None, 3] - bbox[:, None, 1]) / cfg.input_body_shape[0] * cfg.output_hm_shape[1]) /
+                        cfg.output_hand_hm_shape[1])
+                joint_img[:, smpl_x.pos_joint_part[part_name], 1] += (
+                        bbox[:, None, 1] / cfg.input_body_shape[0] * cfg.output_hm_shape[1])
+
+            # change input_body_shape to input_img_shape
+            for bbox in (lhand_bbox, rhand_bbox, face_bbox):
+                bbox[:, 0] *= cfg.input_img_shape[1] / cfg.input_body_shape[1]
+                bbox[:, 1] *= cfg.input_img_shape[0] / cfg.input_body_shape[0]
+                bbox[:, 2] *= cfg.input_img_shape[1] / cfg.input_body_shape[1]
+                bbox[:, 3] *= cfg.input_img_shape[0] / cfg.input_body_shape[0]
+
+            # test output
+            out = {}
+            out['img'] = inputs['img']
+            out['joint_img'] = joint_img
+            out['smplx_joint_proj'] = joint_proj
+            out['smplx_mesh_cam'] = mesh_cam
+            out['smplx_root_pose'] = root_pose
+            out['smplx_body_pose'] = body_pose
+            out['smplx_lhand_pose'] = lhand_pose
+            out['smplx_rhand_pose'] = rhand_pose
+            out['smplx_jaw_pose'] = jaw_pose
+            out['smplx_shape'] = shape
+            out['smplx_expr'] = expr
+            out['cam_trans'] = cam_trans
+            out['lhand_bbox'] = lhand_bbox
+            out['rhand_bbox'] = rhand_bbox
+            out['face_bbox'] = face_bbox
+            if 'smplx_shape' in targets:
+                out['smplx_shape_target'] = targets['smplx_shape']
+            if 'img_path' in meta_info:
+                out['img_path'] = meta_info['img_path']
+            if 'smplx_pose' in targets:
+                out['smplx_mesh_cam_pseudo_gt'] = mesh_pseudo_gt
+            if 'smplx_mesh_cam' in targets:
+                out['smplx_mesh_cam_target'] = targets['smplx_mesh_cam']
+            if 'smpl_mesh_cam' in targets:
+                out['smpl_mesh_cam_target'] = targets['smpl_mesh_cam']
+            if 'bb2img_trans' in meta_info:
+                out['bb2img_trans'] = meta_info['bb2img_trans']
+            if 'gt_smplx_transl' in meta_info:
+                out['gt_smplx_transl'] = meta_info['gt_smplx_transl']
+
+            return out
+
+def init_weights(m):
+    try:
+        if type(m) == nn.ConvTranspose2d:
+            nn.init.normal_(m.weight, std=0.001)
+        elif type(m) == nn.Conv2d:
+            nn.init.normal_(m.weight, std=0.001)
+            nn.init.constant_(m.bias, 0)
+        elif type(m) == nn.BatchNorm2d:
+            nn.init.constant_(m.weight, 1)
+            nn.init.constant_(m.bias, 0)
+        elif type(m) == nn.Linear:
+            nn.init.normal_(m.weight, std=0.01)
+            nn.init.constant_(m.bias, 0)
+    except AttributeError:
+        pass
+
+
+def get_model(mode):
+
+    # body
+    vit_cfg = Config.fromfile(cfg.encoder_config_file)
+    vit = build_posenet(vit_cfg.model)
+    body_position_net = PositionNet('body', feat_dim=cfg.feat_dim)
+    body_rotation_net = BodyRotationNet(feat_dim=cfg.feat_dim)
+    box_net = BoxNet(feat_dim=cfg.feat_dim)
+
+    # hand
+    hand_position_net = PositionNet('hand', feat_dim=cfg.feat_dim)
+    hand_roi_net = HandRoI(feat_dim=cfg.feat_dim, upscale=cfg.upscale)
+    hand_rotation_net = HandRotationNet('hand', feat_dim=cfg.feat_dim)
+
+    # face
+    face_regressor = FaceRegressor(feat_dim=cfg.feat_dim)
+
+    if mode == 'train':
+        # body
+        if not getattr(cfg, 'random_init', False):
+            encoder_pretrained_model = torch.load(cfg.encoder_pretrained_model_path)['state_dict']
+            vit.load_state_dict(encoder_pretrained_model, strict=False)
+            print(f"Initialize encoder from {cfg.encoder_pretrained_model_path}")
+        else:
+            print('Random init!!!!!!!')
+
+        body_position_net.apply(init_weights)
+        body_rotation_net.apply(init_weights)
+        box_net.apply(init_weights)
+
+        # hand
+        hand_position_net.apply(init_weights)
+        hand_roi_net.apply(init_weights)
+        hand_rotation_net.apply(init_weights)
+
+        # face
+        face_regressor.apply(init_weights)
+
+    encoder = vit.backbone
+
+    model = Model(encoder, body_position_net, body_rotation_net, box_net, hand_position_net, hand_roi_net, hand_rotation_net,
+                  face_regressor)
+    return model

main/_base_/datasets/300w.py

@@ -0,0 +1,384 @@
+dataset_info = dict(
+    dataset_name='300w',
+    paper_info=dict(
+        author='Sagonas, Christos and Antonakos, Epameinondas '
+        'and Tzimiropoulos, Georgios and Zafeiriou, Stefanos '
+        'and Pantic, Maja',
+        title='300 faces in-the-wild challenge: '
+        'Database and results',
+        container='Image and vision computing',
+        year='2016',
+        homepage='https://ibug.doc.ic.ac.uk/resources/300-W/',
+    ),
+    keypoint_info={
+        0:
+        dict(
+            name='kpt-0', id=0, color=[255, 255, 255], type='', swap='kpt-16'),
+        1:
+        dict(
+            name='kpt-1', id=1, color=[255, 255, 255], type='', swap='kpt-15'),
+        2:
+        dict(
+            name='kpt-2', id=2, color=[255, 255, 255], type='', swap='kpt-14'),
+        3:
+        dict(
+            name='kpt-3', id=3, color=[255, 255, 255], type='', swap='kpt-13'),
+        4:
+        dict(
+            name='kpt-4', id=4, color=[255, 255, 255], type='', swap='kpt-12'),
+        5:
+        dict(
+            name='kpt-5', id=5, color=[255, 255, 255], type='', swap='kpt-11'),
+        6:
+        dict(
+            name='kpt-6', id=6, color=[255, 255, 255], type='', swap='kpt-10'),
+        7:
+        dict(name='kpt-7', id=7, color=[255, 255, 255], type='', swap='kpt-9'),
+        8:
+        dict(name='kpt-8', id=8, color=[255, 255, 255], type='', swap=''),
+        9:
+        dict(name='kpt-9', id=9, color=[255, 255, 255], type='', swap='kpt-7'),
+        10:
+        dict(
+            name='kpt-10', id=10, color=[255, 255, 255], type='',
+            swap='kpt-6'),
+        11:
+        dict(
+            name='kpt-11', id=11, color=[255, 255, 255], type='',
+            swap='kpt-5'),
+        12:
+        dict(
+            name='kpt-12', id=12, color=[255, 255, 255], type='',
+            swap='kpt-4'),
+        13:
+        dict(
+            name='kpt-13', id=13, color=[255, 255, 255], type='',
+            swap='kpt-3'),
+        14:
+        dict(
+            name='kpt-14', id=14, color=[255, 255, 255], type='',
+            swap='kpt-2'),
+        15:
+        dict(
+            name='kpt-15', id=15, color=[255, 255, 255], type='',
+            swap='kpt-1'),
+        16:
+        dict(
+            name='kpt-16', id=16, color=[255, 255, 255], type='',
+            swap='kpt-0'),
+        17:
+        dict(
+            name='kpt-17',
+            id=17,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-26'),
+        18:
+        dict(
+            name='kpt-18',
+            id=18,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-25'),
+        19:
+        dict(
+            name='kpt-19',
+            id=19,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-24'),
+        20:
+        dict(
+            name='kpt-20',
+            id=20,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-23'),
+        21:
+        dict(
+            name='kpt-21',
+            id=21,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-22'),
+        22:
+        dict(
+            name='kpt-22',
+            id=22,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-21'),
+        23:
+        dict(
+            name='kpt-23',
+            id=23,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-20'),
+        24:
+        dict(
+            name='kpt-24',
+            id=24,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-19'),
+        25:
+        dict(
+            name='kpt-25',
+            id=25,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-18'),
+        26:
+        dict(
+            name='kpt-26',
+            id=26,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-17'),
+        27:
+        dict(name='kpt-27', id=27, color=[255, 255, 255], type='', swap=''),
+        28:
+        dict(name='kpt-28', id=28, color=[255, 255, 255], type='', swap=''),
+        29:
+        dict(name='kpt-29', id=29, color=[255, 255, 255], type='', swap=''),
+        30:
+        dict(name='kpt-30', id=30, color=[255, 255, 255], type='', swap=''),
+        31:
+        dict(
+            name='kpt-31',
+            id=31,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-35'),
+        32:
+        dict(
+            name='kpt-32',
+            id=32,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-34'),
+        33:
+        dict(name='kpt-33', id=33, color=[255, 255, 255], type='', swap=''),
+        34:
+        dict(
+            name='kpt-34',
+            id=34,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-32'),
+        35:
+        dict(
+            name='kpt-35',
+            id=35,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-31'),
+        36:
+        dict(
+            name='kpt-36',
+            id=36,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-45'),
+        37:
+        dict(
+            name='kpt-37',
+            id=37,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-44'),
+        38:
+        dict(
+            name='kpt-38',
+            id=38,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-43'),
+        39:
+        dict(
+            name='kpt-39',
+            id=39,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-42'),
+        40:
+        dict(
+            name='kpt-40',
+            id=40,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-47'),
+        41:
+        dict(
+            name='kpt-41',
+            id=41,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-46'),
+        42:
+        dict(
+            name='kpt-42',
+            id=42,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-39'),
+        43:
+        dict(
+            name='kpt-43',
+            id=43,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-38'),
+        44:
+        dict(
+            name='kpt-44',
+            id=44,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-37'),
+        45:
+        dict(
+            name='kpt-45',
+            id=45,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-36'),
+        46:
+        dict(
+            name='kpt-46',
+            id=46,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-41'),
+        47:
+        dict(
+            name='kpt-47',
+            id=47,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-40'),
+        48:
+        dict(
+            name='kpt-48',
+            id=48,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-54'),
+        49:
+        dict(
+            name='kpt-49',
+            id=49,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-53'),
+        50:
+        dict(
+            name='kpt-50',
+            id=50,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-52'),
+        51:
+        dict(name='kpt-51', id=51, color=[255, 255, 255], type='', swap=''),
+        52:
+        dict(
+            name='kpt-52',
+            id=52,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-50'),
+        53:
+        dict(
+            name='kpt-53',
+            id=53,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-49'),
+        54:
+        dict(
+            name='kpt-54',
+            id=54,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-48'),
+        55:
+        dict(
+            name='kpt-55',
+            id=55,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-59'),
+        56:
+        dict(
+            name='kpt-56',
+            id=56,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-58'),
+        57:
+        dict(name='kpt-57', id=57, color=[255, 255, 255], type='', swap=''),
+        58:
+        dict(
+            name='kpt-58',
+            id=58,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-56'),
+        59:
+        dict(
+            name='kpt-59',
+            id=59,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-55'),
+        60:
+        dict(
+            name='kpt-60',
+            id=60,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-64'),
+        61:
+        dict(
+            name='kpt-61',
+            id=61,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-63'),
+        62:
+        dict(name='kpt-62', id=62, color=[255, 255, 255], type='', swap=''),
+        63:
+        dict(
+            name='kpt-63',
+            id=63,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-61'),
+        64:
+        dict(
+            name='kpt-64',
+            id=64,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-60'),
+        65:
+        dict(
+            name='kpt-65',
+            id=65,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-67'),
+        66:
+        dict(name='kpt-66', id=66, color=[255, 255, 255], type='', swap=''),
+        67:
+        dict(
+            name='kpt-67',
+            id=67,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-65'),
+    },
+    skeleton_info={},
+    joint_weights=[1.] * 68,
+    sigmas=[])

main/_base_/datasets/aflw.py

@@ -0,0 +1,83 @@
+dataset_info = dict(
+    dataset_name='aflw',
+    paper_info=dict(
+        author='Koestinger, Martin and Wohlhart, Paul and '
+        'Roth, Peter M and Bischof, Horst',
+        title='Annotated facial landmarks in the wild: '
+        'A large-scale, real-world database for facial '
+        'landmark localization',
+        container='2011 IEEE international conference on computer '
+        'vision workshops (ICCV workshops)',
+        year='2011',
+        homepage='https://www.tugraz.at/institute/icg/research/'
+        'team-bischof/lrs/downloads/aflw/',
+    ),
+    keypoint_info={
+        0:
+        dict(name='kpt-0', id=0, color=[255, 255, 255], type='', swap='kpt-5'),
+        1:
+        dict(name='kpt-1', id=1, color=[255, 255, 255], type='', swap='kpt-4'),
+        2:
+        dict(name='kpt-2', id=2, color=[255, 255, 255], type='', swap='kpt-3'),
+        3:
+        dict(name='kpt-3', id=3, color=[255, 255, 255], type='', swap='kpt-2'),
+        4:
+        dict(name='kpt-4', id=4, color=[255, 255, 255], type='', swap='kpt-1'),
+        5:
+        dict(name='kpt-5', id=5, color=[255, 255, 255], type='', swap='kpt-0'),
+        6:
+        dict(
+            name='kpt-6', id=6, color=[255, 255, 255], type='', swap='kpt-11'),
+        7:
+        dict(
+            name='kpt-7', id=7, color=[255, 255, 255], type='', swap='kpt-10'),
+        8:
+        dict(name='kpt-8', id=8, color=[255, 255, 255], type='', swap='kpt-9'),
+        9:
+        dict(name='kpt-9', id=9, color=[255, 255, 255], type='', swap='kpt-8'),
+        10:
+        dict(
+            name='kpt-10', id=10, color=[255, 255, 255], type='',
+            swap='kpt-7'),
+        11:
+        dict(
+            name='kpt-11', id=11, color=[255, 255, 255], type='',
+            swap='kpt-6'),
+        12:
+        dict(
+            name='kpt-12',
+            id=12,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-14'),
+        13:
+        dict(name='kpt-13', id=13, color=[255, 255, 255], type='', swap=''),
+        14:
+        dict(
+            name='kpt-14',
+            id=14,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-12'),
+        15:
+        dict(
+            name='kpt-15',
+            id=15,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-17'),
+        16:
+        dict(name='kpt-16', id=16, color=[255, 255, 255], type='', swap=''),
+        17:
+        dict(
+            name='kpt-17',
+            id=17,
+            color=[255, 255, 255],
+            type='',
+            swap='kpt-15'),
+        18:
+        dict(name='kpt-18', id=18, color=[255, 255, 255], type='', swap='')
+    },
+    skeleton_info={},
+    joint_weights=[1.] * 19,
+    sigmas=[])

main/_base_/datasets/aic.py

@@ -0,0 +1,140 @@
+dataset_info = dict(
+    dataset_name='aic',
+    paper_info=dict(
+        author='Wu, Jiahong and Zheng, He and Zhao, Bo and '
+        'Li, Yixin and Yan, Baoming and Liang, Rui and '
+        'Wang, Wenjia and Zhou, Shipei and Lin, Guosen and '
+        'Fu, Yanwei and others',
+        title='Ai challenger: A large-scale dataset for going '
+        'deeper in image understanding',
+        container='arXiv',
+        year='2017',
+        homepage='https://github.com/AIChallenger/AI_Challenger_2017',
+    ),
+    keypoint_info={
+        0:
+        dict(
+            name='right_shoulder',
+            id=0,
+            color=[255, 128, 0],
+            type='upper',
+            swap='left_shoulder'),
+        1:
+        dict(
+            name='right_elbow',
+            id=1,
+            color=[255, 128, 0],
+            type='upper',
+            swap='left_elbow'),
+        2:
+        dict(
+            name='right_wrist',
+            id=2,
+            color=[255, 128, 0],
+            type='upper',
+            swap='left_wrist'),
+        3:
+        dict(
+            name='left_shoulder',
+            id=3,
+            color=[0, 255, 0],
+            type='upper',
+            swap='right_shoulder'),
+        4:
+        dict(
+            name='left_elbow',
+            id=4,
+            color=[0, 255, 0],
+            type='upper',
+            swap='right_elbow'),
+        5:
+        dict(
+            name='left_wrist',
+            id=5,
+            color=[0, 255, 0],
+            type='upper',
+            swap='right_wrist'),
+        6:
+        dict(
+            name='right_hip',
+            id=6,
+            color=[255, 128, 0],
+            type='lower',
+            swap='left_hip'),
+        7:
+        dict(
+            name='right_knee',
+            id=7,
+            color=[255, 128, 0],
+            type='lower',
+            swap='left_knee'),
+        8:
+        dict(
+            name='right_ankle',
+            id=8,
+            color=[255, 128, 0],
+            type='lower',
+            swap='left_ankle'),
+        9:
+        dict(
+            name='left_hip',
+            id=9,
+            color=[0, 255, 0],
+            type='lower',
+            swap='right_hip'),
+        10:
+        dict(
+            name='left_knee',
+            id=10,
+            color=[0, 255, 0],
+            type='lower',
+            swap='right_knee'),
+        11:
+        dict(
+            name='left_ankle',
+            id=11,
+            color=[0, 255, 0],
+            type='lower',
+            swap='right_ankle'),
+        12:
+        dict(
+            name='head_top',
+            id=12,
+            color=[51, 153, 255],
+            type='upper',
+            swap=''),
+        13:
+        dict(name='neck', id=13, color=[51, 153, 255], type='upper', swap='')
+    },
+    skeleton_info={
+        0:
+        dict(link=('right_wrist', 'right_elbow'), id=0, color=[255, 128, 0]),
+        1: dict(
+            link=('right_elbow', 'right_shoulder'), id=1, color=[255, 128, 0]),
+        2: dict(link=('right_shoulder', 'neck'), id=2, color=[51, 153, 255]),
+        3: dict(link=('neck', 'left_shoulder'), id=3, color=[51, 153, 255]),
+        4: dict(link=('left_shoulder', 'left_elbow'), id=4, color=[0, 255, 0]),
+        5: dict(link=('left_elbow', 'left_wrist'), id=5, color=[0, 255, 0]),
+        6: dict(link=('right_ankle', 'right_knee'), id=6, color=[255, 128, 0]),
+        7: dict(link=('right_knee', 'right_hip'), id=7, color=[255, 128, 0]),
+        8: dict(link=('right_hip', 'left_hip'), id=8, color=[51, 153, 255]),
+        9: dict(link=('left_hip', 'left_knee'), id=9, color=[0, 255, 0]),
+        10: dict(link=('left_knee', 'left_ankle'), id=10, color=[0, 255, 0]),
+        11: dict(link=('head_top', 'neck'), id=11, color=[51, 153, 255]),
+        12: dict(
+            link=('right_shoulder', 'right_hip'), id=12, color=[51, 153, 255]),
+        13:
+        dict(link=('left_shoulder', 'left_hip'), id=13, color=[51, 153, 255])
+    },
+    joint_weights=[
+        1., 1.2, 1.5, 1., 1.2, 1.5, 1., 1.2, 1.5, 1., 1.2, 1.5, 1., 1.
+    ],
+
+    # 'https://github.com/AIChallenger/AI_Challenger_2017/blob/master/'
+    # 'Evaluation/keypoint_eval/keypoint_eval.py#L50'
+    # delta = 2 x sigma
+    sigmas=[
+        0.01388152, 0.01515228, 0.01057665, 0.01417709, 0.01497891, 0.01402144,
+        0.03909642, 0.03686941, 0.01981803, 0.03843971, 0.03412318, 0.02415081,
+        0.01291456, 0.01236173
+    ])