init

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+body_models
+results
+weights
+tada-extend

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 xin he
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

SMPLX/config_files/smpl2smplh.yaml

@@ -0,0 +1,25 @@
+datasets:
+    mesh_folder:
+        data_folder: 'transfer_data/meshes/smpl'
+deformation_transfer_path: 'transfer_data/smpl2smplh_def_transfer.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'trust-ncg'
+    maxiters: 100
+    gtol: 1e-06
+
+body_model:
+    model_type: "smplh"
+    # SMPL+H has no neutral model, so we have to manually select the gender
+    gender: "female"
+    # gender: "male"
+    folder: "transfer_data/body_models"
+    use_compressed: False
+    smplh:
+        betas:
+            num: 10

SMPLX/config_files/smpl2smplx.yaml

@@ -0,0 +1,26 @@
+datasets:
+    mesh_folder:
+        data_folder: 'transfer_data/meshes/smpl'
+deformation_transfer_path: 'transfer_data/smpl2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'trust-ncg'
+    maxiters: 100
+    gtol: 1e-06
+
+body_model:
+    model_type: "smplx"
+    gender: "neutral"
+    folder: "transfer_data/body_models"
+    use_compressed: False
+    use_face_contour: True
+    smplx:
+        betas:
+            num: 10
+        expression:
+            num: 10

SMPLX/config_files/smplh2smpl.yaml

@@ -0,0 +1,24 @@
+datasets:
+    mesh_folder:
+        data_folder: 'transfer_data/meshes/smplh'
+deformation_transfer_path: 'transfer_data/smplh2smpl_def_transfer.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'trust-ncg'
+    maxiters: 100
+    gtol: 1e-06
+
+body_model:
+    model_type: "smpl"
+    gender: "neutral"
+    folder: "transfer_data/body_models"
+    use_compressed: False
+    use_face_contour: True
+    smpl:
+        betas:
+            num: 10

SMPLX/config_files/smplh2smplx.yaml

@@ -0,0 +1,26 @@
+datasets:
+    mesh_folder:
+        data_folder: 'transfer_data/meshes/smplh'
+deformation_transfer_path: 'transfer_data/smplh2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'trust-ncg'
+    maxiters: 100
+    gtol: 1e-06
+
+body_model:
+    model_type: "smplx"
+    gender: "neutral"
+    folder: "transfer_data/body_models"
+    use_compressed: False
+    use_face_contour: True
+    smplx:
+        betas:
+            num: 10
+        expression:
+            num: 10

SMPLX/config_files/smplh2smplx_as.yaml

@@ -0,0 +1,26 @@
+datasets:
+    mesh_folder:
+        data_folder: 'transfer_data/meshes/amass_sample'
+deformation_transfer_path: 'transfer_data/smplh2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'trust-ncg'
+    maxiters: 100
+    gtol: 1e-06
+
+body_model:
+    model_type: "smplx"
+    gender: "male"
+    folder: "/data/TTA/data/body_models"
+    use_compressed: False
+    use_face_contour: True
+    smplx:
+        betas:
+            num: 10
+        expression:
+            num: 10

SMPLX/config_files/smplh2smplx_onepose.yaml

@@ -0,0 +1,27 @@
+datasets:
+    mesh_folder:
+        data_folder: 'transfer_data/meshes/amass_onepose'
+deformation_transfer_path: 'transfer_data/smplh2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'adam'
+    lr: 0.1
+    maxiters: 10000
+    gtol: 1e-06
+
+body_model:
+    model_type: "smplx"
+    gender: "neutral"
+    folder: "models"
+    use_compressed: False
+    use_face_contour: True
+    smplx:
+        betas:
+            num: 10
+        expression:
+            num: 10

SMPLX/config_files/smplx2smpl.yaml

@@ -0,0 +1,25 @@
+datasets:
+    mesh_folder:
+        data_folder: 'meshes/smplx'
+deformation_transfer_path: 'transfer_data/smplx2smpl_deftrafo_setup.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'lbfgs'
+    maxiters: 200
+    gtol: 1e-06
+
+body_model:
+    model_type: "smpl"
+    gender: "neutral"
+    ext: 'pkl'
+    folder: "transfer_data/body_models"
+    use_compressed: False
+    use_face_contour: True
+    smpl:
+        betas:
+            num: 10

SMPLX/config_files/smplx2smplh.yaml

@@ -0,0 +1,27 @@
+datasets:
+    mesh_folder:
+        data_folder: 'meshes/smplx'
+deformation_transfer_path: 'transfer_data/smplx2smplh_deftrafo_setup.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+    per_part: False
+
+optim:
+    type: 'lbfgs'
+    maxiters: 200
+    gtol: 1e-06
+
+body_model:
+    model_type: "smplh"
+    # SMPL+H has no neutral model, so we have to manually select the gender
+    gender: "female"
+    # gender: "male"
+    ext: 'pkl'
+    folder: "transfer_data/body_models"
+    use_compressed: False
+    use_face_contour: True
+    smplh:
+        betas:
+            num: 10

SMPLX/joints2smpl.py

@@ -0,0 +1,59 @@
+import torch
+from SMPLX.visualize_joint2smpl.simplify_loc2rot import joints2smpl
+import argparse
+import numpy as np
+import os
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description='transfer joint3d to smpls')
+parser.add_argument("--model_path", default="/data/TTA/data/body_models")
+parser.add_argument('--source_path', default="/data/TTA/data/humanact12/group_000")
+parser.add_argument("--target_path", default="/data/TTA/data/humanact_smplh/group_000")
+parser.add_argument("--mode", default="joints", choices=["t2m", "joints"])
+args = parser.parse_args()
+device = "cuda"
+
+if os.path.isdir(args.source_path):
+    os.makedirs(args.target_path, exist_ok=True)
+    files = os.listdir(args.source_path)
+    target_files = files
+else:
+    files = [args.source_path]
+    args.source_path = ""
+
+    if args.target_path.split(".")[-1] != "npy":
+        os.makedirs(args.target_path)
+        target_files = [files[0].split("/")[-1]]
+    else:
+        target_files = [args.target_path]
+        args.target_path = ""
+
+for i in range(len(files)):
+    curr_path = os.path.join(args.source_path, files[i])
+    target_path = os.path.join(args.target_path, target_files[i])
+    if os.path.exists(target_path):
+        continue
+
+    curr_file = np.load(curr_path)       #### [nframe, 263]
+    curr_file = torch.from_numpy(curr_file)
+
+    if args.mode == "t2m":
+        from dataset.t2m.recover_joints import recover_from_ric
+        motions = recover_from_ric(curr_file, 22)    #### [nframes, 22, 3]
+        motions = motions.detach().cpu().numpy()
+    else:
+        motions = curr_file.detach().cpu().numpy()
+ 
+    frames, njoints, nfeats = motions.shape
+    MINS = motions.min(axis=0).min(axis=0)
+    MAXS = motions.max(axis=0).max(axis=0)
+    height_offset = MINS[1]
+    motions[:, :, 1] -= height_offset
+    model = joints2smpl(frames, 0, True, model_path=args.model_path)
+    target, trans = model.joint2smpl(motions)
+
+    target = np.concatenate([target, trans], axis=1)
+
+    np.save(target_path, target)
+    if i % 10 == 0:
+        print("save %d npys"%(i))

SMPLX/read_from_npy.py

@@ -0,0 +1,108 @@
+import numpy as np
+import torch
+
+def npy2info(motions, num_shapes=10):
+    if isinstance(motions, str):
+        motions = np.load(motions)
+        
+    trans = None
+    gnum = 2
+
+    if isinstance(motions, np.ndarray):
+        betas = np.zeros([motions.shape[0], num_shapes]).astype(motions.dtype)
+    else:
+        betas = torch.zeros([motions.shape[0], num_shapes], dtype=motions.dtype)
+
+    if len(motions.shape) == 3:
+        motions = motions.reshape(motions.shape[0], -1)
+
+    if motions.shape[1] in [73, 157, 166]:
+        gnum = motions[:, -1:][0]
+        motions = motions[:, :-1]
+    elif motions.shape[1] in [75, 159, 168]:
+        gnum = 2
+        trans = motions[:, -3::]
+        motions = motions[:, :-3]
+    elif motions.shape[1] in [76, 160, 169]:
+        gnum = motions[:, -1:][0]
+        trans = motions[:, -4:-1:]
+        motions = motions[:, :-4]     
+    elif motions.shape[1] in [72 + num_shapes, 156 + num_shapes, 165 + num_shapes]:
+        betas = motions[:, -num_shapes::]
+        gnum = 2
+        motions = motions[:, :-num_shapes]
+    elif motions.shape[1] in [73 + num_shapes, 157 + num_shapes, 166 + num_shapes]:
+        betas = motions[:, -num_shapes::]
+        gnum = motions[:, -num_shapes-1:-num_shapes:][0]
+        motions = motions[:, :-num_shapes-1]      
+    elif motions.shape[1] in [75 + num_shapes, 159 + num_shapes, 168 + num_shapes]:
+        betas = motions[:, -num_shapes::]
+        gnum = 2
+        trans = motions[:, -num_shapes-3:-num_shapes:]
+        motions = motions[:, :-num_shapes-3]      
+    elif motions.shape[1] in [76 + num_shapes, 160 + num_shapes, 169 + num_shapes]:
+        betas = motions[:, -num_shapes::]
+        gnum = motions[:, -num_shapes-1:-num_shapes:][0]
+        trans = motions[:, -num_shapes-4:-num_shapes-1:]
+        motions = motions[:, :-num_shapes-4]    
+
+    if gnum == 0:
+        gender = "female"
+    elif gnum == 1:
+        gender = "male"
+    else:
+        gender = "neutral"
+
+    return motions, trans, gender, betas
+
+def info2dict(pose, trans=None, betas=None, mode="smpl", device="cuda", index=-1):
+    if isinstance(pose, np.ndarray):
+        pose = torch.from_numpy(pose)
+
+    if trans is not None and isinstance(trans, np.ndarray):
+        trans = torch.from_numpy(trans)
+    
+    if betas is not None and isinstance(betas, np.ndarray):
+        betas = torch.from_numpy(betas)
+    elif betas is None:
+        betas = torch.zeros([pose.shape[0], 10])
+
+    if index != -1:
+        pose = pose[index:index+1]
+
+        if trans is not None:
+            trans = trans[index:index+1]
+
+        betas = betas[index:index+1]
+
+    if mode == "smplx":
+        inputs = {
+            "global_orient": pose[:, :3].float().to(device),
+            "body_pose": pose[:, 3:66].float().to(device),
+            "jaw_pose": pose[:, 66:69].float().to(device),
+            "leye_pose": pose[:, 69:72].float().to(device),
+            "reye_pose": pose[:, 72:75].float().to(device),
+            "left_hand_pose":pose[:, 75:120].float().to(device),
+            "right_hand_pose":pose[:, 120:].float().to(device),
+        }
+    elif mode == "smplh":
+        inputs = {
+            "global_orient": pose[:, :3].float().to(device),
+            "body_pose": pose[:, 3:66].float().to(device),
+            "left_hand_pose":pose[:, 66:111].float().to(device),
+            "right_hand_pose":pose[:, 111:].float().to(device),
+        } 
+    elif mode == "smpl":
+        inputs = {
+            "global_orient": pose[:, :3].float().to(device),
+            "body_pose": pose[:, 3:].float().to(device),
+        }      
+
+    if trans is not None:
+        inputs["transl"] = trans[:, :].float().to(device)
+    else:
+        print("No Translation Information")
+
+    inputs["betas"] = betas[:, :].float().to(device)
+
+    return inputs

SMPLX/read_joints_from_pose.py

@@ -0,0 +1,110 @@
+import torch
+import numpy as np
+from torch import nn
+import pickle as pkl
+import torch.nn.functional as F
+
+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)
+
+
+class Get_Joints(nn.Module):
+    def __init__(self, path, batch_size=300) -> None:
+        super().__init__()
+        self.betas = nn.parameter.Parameter(torch.zeros([batch_size, 10], dtype=torch.float32), requires_grad=False)
+        with open(path, "rb") as f:
+            smpl_prior = pkl.load(f, encoding="latin1")
+            data_struct = Struct(**smpl_prior)
+
+        self.v_template = nn.parameter.Parameter(torch.from_numpy(to_np(data_struct.v_template)), requires_grad=False)
+        self.shapedirs = nn.parameter.Parameter(torch.from_numpy(to_np(data_struct.shapedirs)), requires_grad=False)
+        self.J_regressor = nn.parameter.Parameter(torch.from_numpy(to_np(data_struct.J_regressor)), requires_grad=False)
+        posedirs = torch.from_numpy(to_np(data_struct.posedirs))
+        num_pose_basis = posedirs.shape[-1]
+        posedirs = posedirs.reshape([-1, num_pose_basis]).permute(1, 0)
+        self.posedirs = nn.parameter.Parameter(posedirs, requires_grad=False)
+        self.parents = nn.parameter.Parameter(torch.from_numpy(to_np(data_struct.kintree_table)[0]).long(), requires_grad=False)
+        self.parents[0] = -1
+
+        self.ident = nn.parameter.Parameter(torch.eye(3), requires_grad=False)
+        self.K = nn.parameter.Parameter(torch.zeros([1, 3, 3]), requires_grad=False)
+        self.zeros = nn.parameter.Parameter(torch.zeros([1, 1]), requires_grad=False)
+
+    def blend_shapes(self, betas, shape_disps):
+        blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
+        return blend_shape
+
+    def vertices2joints(self, J_regressor, vertices):
+        return torch.einsum('bik,ji->bjk', [vertices, J_regressor])
+
+    def batch_rodrigues(
+        self,
+        rot_vecs,
+        epsilon = 1e-8,
+    ):
+        batch_size = rot_vecs.shape[0]
+        angle = torch.norm(rot_vecs + epsilon, 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 = self.K.repeat(batch_size, 1, 1)
+        zeros = self.zeros.repeat(batch_size, 1)
+        K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1).view((batch_size, 3, 3))
+        ident = self.ident.unsqueeze(0) 
+        rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
+        return rot_mat
+
+    def transform_mat(self, R, t):
+        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(
+        self,
+        rot_mats,
+        joints,
+        parents,
+    ):
+        joints = torch.unsqueeze(joints, dim=-1)
+
+        rel_joints = joints.clone()
+        rel_joints[:, 1:] -= joints[:, parents[1:]]
+
+        transforms_mat = self.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]
+        return posed_joints
+
+    def forward(self, pose, trans=None):
+        pose = pose.float()
+        batch = pose.shape[0]
+        betas = self.betas[:batch]
+        v_shaped = self.v_template + self.blend_shapes(betas, self.shapedirs)
+        J = self.vertices2joints(self.J_regressor, v_shaped)
+        rot_mats = self.batch_rodrigues(pose.view(-1, 3)).view([batch, -1, 3, 3])
+        J_transformed = self.batch_rigid_transform(rot_mats, J, self.parents)
+        if trans is not None:
+            J_transformed += trans.unsqueeze(dim=1)
+        return J_transformed

SMPLX/rotation_conversions.py

@@ -0,0 +1,532 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+# Check PYTORCH3D_LICENCE before use
+
+import functools
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+
+"""
+The transformation matrices returned from the functions in this file assume
+the points on which the transformation will be applied are column vectors.
+i.e. the R matrix is structured as
+    R = [
+            [Rxx, Rxy, Rxz],
+            [Ryx, Ryy, Ryz],
+            [Rzx, Rzy, Rzz],
+        ]  # (3, 3)
+This matrix can be applied to column vectors by post multiplication
+by the points e.g.
+    points = [[0], [1], [2]]  # (3 x 1) xyz coordinates of a point
+    transformed_points = R * points
+To apply the same matrix to points which are row vectors, the R matrix
+can be transposed and pre multiplied by the points:
+e.g.
+    points = [[0, 1, 2]]  # (1 x 3) xyz coordinates of a point
+    transformed_points = points * R.transpose(1, 0)
+"""
+
+
+def quaternion_to_matrix(quaternions):
+    """
+    Convert rotations given as quaternions to rotation matrices.
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def _copysign(a, b):
+    """
+    Return a tensor where each element has the absolute value taken from the,
+    corresponding element of a, with sign taken from the corresponding
+    element of b. This is like the standard copysign floating-point operation,
+    but is not careful about negative 0 and NaN.
+    Args:
+        a: source tensor.
+        b: tensor whose signs will be used, of the same shape as a.
+    Returns:
+        Tensor of the same shape as a with the signs of b.
+    """
+    signs_differ = (a < 0) != (b < 0)
+    return torch.where(signs_differ, -a, a)
+
+
+def _sqrt_positive_part(x):
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    ret[positive_mask] = torch.sqrt(x[positive_mask])
+    return ret
+
+
+def matrix_to_quaternion(matrix):
+    """
+    Convert rotations given as rotation matrices to quaternions.
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix  shape f{matrix.shape}.")
+    m00 = matrix[..., 0, 0]
+    m11 = matrix[..., 1, 1]
+    m22 = matrix[..., 2, 2]
+    o0 = 0.5 * _sqrt_positive_part(1 + m00 + m11 + m22)
+    x = 0.5 * _sqrt_positive_part(1 + m00 - m11 - m22)
+    y = 0.5 * _sqrt_positive_part(1 - m00 + m11 - m22)
+    z = 0.5 * _sqrt_positive_part(1 - m00 - m11 + m22)
+    o1 = _copysign(x, matrix[..., 2, 1] - matrix[..., 1, 2])
+    o2 = _copysign(y, matrix[..., 0, 2] - matrix[..., 2, 0])
+    o3 = _copysign(z, matrix[..., 1, 0] - matrix[..., 0, 1])
+    return torch.stack((o0, o1, o2, o3), -1)
+
+
+def _axis_angle_rotation(axis: str, angle):
+    """
+    Return the rotation matrices for one of the rotations about an axis
+    of which Euler angles describe, for each value of the angle given.
+    Args:
+        axis: Axis label "X" or "Y or "Z".
+        angle: any shape tensor of Euler angles in radians
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+
+    cos = torch.cos(angle)
+    sin = torch.sin(angle)
+    one = torch.ones_like(angle)
+    zero = torch.zeros_like(angle)
+
+    if axis == "X":
+        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
+    if axis == "Y":
+        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
+    if axis == "Z":
+        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
+
+    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
+
+
+def euler_angles_to_matrix(euler_angles, convention: str):
+    """
+    Convert rotations given as Euler angles in radians to rotation matrices.
+    Args:
+        euler_angles: Euler angles in radians as tensor of shape (..., 3).
+        convention: Convention string of three uppercase letters from
+            {"X", "Y", and "Z"}.
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
+        raise ValueError("Invalid input euler angles.")
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    matrices = map(_axis_angle_rotation, convention, torch.unbind(euler_angles, -1))
+    return functools.reduce(torch.matmul, matrices)
+
+
+def _angle_from_tan(
+    axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
+):
+    """
+    Extract the first or third Euler angle from the two members of
+    the matrix which are positive constant times its sine and cosine.
+    Args:
+        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
+        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
+            convention.
+        data: Rotation matrices as tensor of shape (..., 3, 3).
+        horizontal: Whether we are looking for the angle for the third axis,
+            which means the relevant entries are in the same row of the
+            rotation matrix. If not, they are in the same column.
+        tait_bryan: Whether the first and third axes in the convention differ.
+    Returns:
+        Euler Angles in radians for each matrix in data as a tensor
+        of shape (...).
+    """
+
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return torch.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return torch.atan2(-data[..., i2], data[..., i1])
+    return torch.atan2(data[..., i2], -data[..., i1])
+
+
+def _index_from_letter(letter: str):
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+
+
+def matrix_to_euler_angles(matrix, convention: str):
+    """
+    Convert rotations given as rotation matrices to Euler angles in radians.
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+        convention: Convention string of three uppercase letters.
+    Returns:
+        Euler angles in radians as tensor of shape (..., 3).
+    """
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix  shape f{matrix.shape}.")
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+    if tait_bryan:
+        central_angle = torch.asin(
+            matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
+        )
+    else:
+        central_angle = torch.acos(matrix[..., i0, i0])
+
+    o = (
+        _angle_from_tan(
+            convention[0], convention[1], matrix[..., i2], False, tait_bryan
+        ),
+        central_angle,
+        _angle_from_tan(
+            convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
+        ),
+    )
+    return torch.stack(o, -1)
+
+
+def random_quaternions(
+    n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate random quaternions representing rotations,
+    i.e. versors with nonnegative real part.
+    Args:
+        n: Number of quaternions in a batch to return.
+        dtype: Type to return.
+        device: Desired device of returned tensor. Default:
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+    Returns:
+        Quaternions as tensor of shape (N, 4).
+    """
+    o = torch.randn((n, 4), dtype=dtype, device=device, requires_grad=requires_grad)
+    s = (o * o).sum(1)
+    o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
+    return o
+
+
+def random_rotations(
+    n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate random rotations as 3x3 rotation matrices.
+    Args:
+        n: Number of rotation matrices in a batch to return.
+        dtype: Type to return.
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+    Returns:
+        Rotation matrices as tensor of shape (n, 3, 3).
+    """
+    quaternions = random_quaternions(
+        n, dtype=dtype, device=device, requires_grad=requires_grad
+    )
+    return quaternion_to_matrix(quaternions)
+
+
+def random_rotation(
+    dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate a single random 3x3 rotation matrix.
+    Args:
+        dtype: Type to return
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set
+    Returns:
+        Rotation matrix as tensor of shape (3, 3).
+    """
+    return random_rotations(1, dtype, device, requires_grad)[0]
+
+
+def standardize_quaternion(quaternions):
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+
+def quaternion_raw_multiply(a, b):
+    """
+    Multiply two quaternions.
+    Usual torch rules for broadcasting apply.
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+    Returns:
+        The product of a and b, a tensor of quaternions shape (..., 4).
+    """
+    aw, ax, ay, az = torch.unbind(a, -1)
+    bw, bx, by, bz = torch.unbind(b, -1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return torch.stack((ow, ox, oy, oz), -1)
+
+
+def quaternion_multiply(a, b):
+    """
+    Multiply two quaternions representing rotations, returning the quaternion
+    representing their composition, i.e. the versor with nonnegative real part.
+    Usual torch rules for broadcasting apply.
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+    Returns:
+        The product of a and b, a tensor of quaternions of shape (..., 4).
+    """
+    ab = quaternion_raw_multiply(a, b)
+    return standardize_quaternion(ab)
+
+
+def quaternion_invert(quaternion):
+    """
+    Given a quaternion representing rotation, get the quaternion representing
+    its inverse.
+    Args:
+        quaternion: Quaternions as tensor of shape (..., 4), with real part
+            first, which must be versors (unit quaternions).
+    Returns:
+        The inverse, a tensor of quaternions of shape (..., 4).
+    """
+
+    return quaternion * quaternion.new_tensor([1, -1, -1, -1])
+
+
+def quaternion_apply(quaternion, point):
+    """
+    Apply the rotation given by a quaternion to a 3D point.
+    Usual torch rules for broadcasting apply.
+    Args:
+        quaternion: Tensor of quaternions, real part first, of shape (..., 4).
+        point: Tensor of 3D points of shape (..., 3).
+    Returns:
+        Tensor of rotated points of shape (..., 3).
+    """
+    if point.size(-1) != 3:
+        raise ValueError(f"Points are not in 3D, f{point.shape}.")
+    real_parts = point.new_zeros(point.shape[:-1] + (1,))
+    point_as_quaternion = torch.cat((real_parts, point), -1)
+    out = quaternion_raw_multiply(
+        quaternion_raw_multiply(quaternion, point_as_quaternion),
+        quaternion_invert(quaternion),
+    )
+    return out[..., 1:]
+
+
+def axis_angle_to_matrix(axis_angle):
+    """
+    Convert rotations given as axis/angle to rotation matrices.
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
+def matrix_to_axis_angle(matrix):
+    """
+    Convert rotations given as rotation matrices to axis/angle.
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
+
+
+def axis_angle_to_quaternion(axis_angle):
+    """
+    Convert rotations given as axis/angle to quaternions.
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = 0.5 * angles
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    quaternions = torch.cat(
+        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
+    )
+    return quaternions
+
+
+def quaternion_to_axis_angle(quaternions):
+    """
+    Convert rotations given as quaternions to axis/angle.
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
+    half_angles = torch.atan2(norms, quaternions[..., :1])
+    angles = 2 * half_angles
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    return quaternions[..., 1:] / sin_half_angles_over_angles
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    """
+    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+    using Gram--Schmidt orthogonalisation per Section B of [1].
+    Args:
+        d6: 6D rotation representation, of size (*, 6)
+    Returns:
+        batch of rotation matrices of size (*, 3, 3)
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1 = F.normalize(a1, dim=-1)
+    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+    b2 = F.normalize(b2, dim=-1)
+    b3 = torch.cross(b1, b2, dim=-1)
+    return torch.stack((b1, b2, b3), dim=-2)
+
+
+def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
+    by dropping the last row. Note that 6D representation is not unique.
+    Args:
+        matrix: batch of rotation matrices of size (*, 3, 3)
+    Returns:
+        6D rotation representation, of size (*, 6)
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+    return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)
+
+def canonicalize_smplh(poses, trans = None):
+    bs, nframes, njoints = poses.shape[:3]
+
+    global_orient = poses[:, :, 0]
+
+    # first global rotations
+    rot2d = matrix_to_axis_angle(global_orient[:, 0])
+    #rot2d[:, :2] = 0  # Remove the rotation along the vertical axis
+    rot2d = axis_angle_to_matrix(rot2d)
+
+    # Rotate the global rotation to eliminate Z rotations
+    global_orient = torch.einsum("ikj,imkl->imjl", rot2d, global_orient)
+
+    # Construct canonicalized version of x
+    xc = torch.cat((global_orient[:, :, None], poses[:, :, 1:]), dim=2)
+
+    if trans is not None:
+        vel = trans[:, 1:] - trans[:, :-1]
+        # Turn the translation as well
+        vel = torch.einsum("ikj,ilk->ilj", rot2d, vel)
+        trans = torch.cat((torch.zeros(bs, 1, 3, device=vel.device),
+                           torch.cumsum(vel, 1)), 1)
+        return xc, trans
+    else:
+        return xc
+    
+    

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: ps-license@tuebingen.mpg.de
+
+from .body_models import (
+    create,
+    SMPL,
+    SMPLH,
+    SMPLX,
+    MANO,
+    FLAME,
+    build_layer,
+    SMPLLayer,
+    SMPLHLayer,
+    SMPLXLayer,
+    MANOLayer,
+    FLAMELayer,
+)

SMPLX/smplx/joint_names.py

@@ -0,0 +1,320 @@
+# -*- 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: ps-license@tuebingen.mpg.de
+
+import numpy as np
+
+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",
+]
+
+
+SMPLH_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",
+    "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",
+]
+
+SMPL_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",
+    "left_hand",
+    "right_hand",
+]
+
+
+class Body:
+    """
+    Class for storing a single body pose.
+    """
+
+    def __init__(self, joints, joint_names):
+        assert joints.ndim > 1
+        assert joints.shape[0] == len(joint_names)
+        self.joints = {}
+        for i, j in enumerate(joint_names):
+            self.joints[j] = joints[i]
+
+    @staticmethod
+    def from_smpl(joints):
+        """
+        Create a Body object from SMPL joints.
+        """
+        return Body(joints, SMPL_JOINT_NAMES)
+
+    @staticmethod
+    def from_smplh(joints):
+        """
+        Create a Body object from SMPLH joints.
+        """
+        return Body(joints, SMPLH_JOINT_NAMES)
+
+    def _as(self, joint_names):
+        """
+        Return a Body object with the specified joint names.
+        """
+        joint_list = []
+        for j in joint_names:
+            if j not in self.joints:
+                joint_list.append(np.zeros_like(self.joints["spine1"]))
+            else:
+                joint_list.append(self.joints[j])
+        return np.stack(joint_list, axis=0)
+
+    def as_smpl(self):
+        """
+        Convert the body to SMPL joints.
+        """
+        return self._as(SMPL_JOINT_NAMES)
+
+    def as_smplh(self):
+        """
+        Convert the body to SMPLH joints.
+        """
+        return self._as(SMPLH_JOINT_NAMES)

SMPLX/smplx/lbs.py

@@ -0,0 +1,405 @@
+# -*- 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: ps-license@tuebingen.mpg.de
+
+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 (soubhik.sanyal@tuebingen.mpg.de)
+        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).to(torch.long)).view(
+        batch_size, -1, 3)
+                        #The '.to(torch.long)'.
+                        # added to make the trace work in c++,
+                        # otherwise you get a runtime error in c++:
+                        # 'index_select(): Expected dtype int32 or int64 for index'
+
+    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]
+
+    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