feature: adding files for rendering

.gitignore

@@ -1,2 +1,3 @@
 /.DS_STORE
-/.env
+/.env
+output/

body_models/smpl/J_regressor_extra.npy

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

body_models/smpl/SMPL_NEUTRAL.pkl

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

body_models/smpl/kintree_table.pkl

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

body_models/smpl/smplfaces.npy

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

dataset/prepare/download_smpl.sh

@@ -0,0 +1,13 @@
+
+mkdir -p body_models
+cd body_models/
+
+echo -e "The smpl files will be stored in the 'body_models/smpl/' folder\n"
+gdown 1INYlGA76ak_cKGzvpOV2Pe6RkYTlXTW2
+rm -rf smpl
+
+unzip smpl.zip
+echo -e "Cleaning\n"
+rm smpl.zip
+
+echo -e "Downloading done!"

generate.py

@@ -6,6 +6,7 @@ import models.t2m_trans as trans
 import models.t2m_trans_uplow as trans_uplow
 import numpy as np
 from exit.utils import visualize_2motions
+from exit.utils import recover_from_ric
 import options.option_transformer as option_trans
 
 
@@ -288,6 +289,14 @@ if __name__ == '__main__':
 
     mmm = MMM(args)
     pred_pose = mmm([args.text], torch.tensor([args.length]), rand_pos=False)
+    num_joints = 22
+    
+    pred_pose = pred_pose[:args.length, :].detach().cpu()
+    
+    converted_pose = recover_from_ric(pred_pose[0].detach().cpu(), num_joints).unsqueeze(0).numpy()
+    
+    np.save('./output/mmm-pred.npy', converted_pose)
+    print('File saved successfully')
 
     std = np.load('./exit/t2m-std.npy')
     mean = np.load('./exit/t2m-mean.npy')

models/rotation2xyz.py

@@ -0,0 +1,92 @@
+# This code is based on https://github.com/Mathux/ACTOR.git
+import torch
+import utils.rotation_conversions as geometry
+
+
+from models.smpl import SMPL, JOINTSTYPE_ROOT
+# from .get_model import JOINTSTYPES
+JOINTSTYPES = ["a2m", "a2mpl", "smpl", "vibe", "vertices"]
+
+
+class Rotation2xyz:
+    def __init__(self, device, dataset='amass'):
+        self.device = device
+        self.dataset = dataset
+        self.smpl_model = SMPL().eval().to(device)
+
+    def __call__(self, x, mask, pose_rep, translation, glob,
+                 jointstype, vertstrans, betas=None, beta=0,
+                 glob_rot=None, get_rotations_back=False, **kwargs):
+        if pose_rep == "xyz":
+            return x
+
+        if mask is None:
+            mask = torch.ones((x.shape[0], x.shape[-1]), dtype=bool, device=x.device)
+
+        if not glob and glob_rot is None:
+            raise TypeError("You must specify global rotation if glob is False")
+
+        if jointstype not in JOINTSTYPES:
+            raise NotImplementedError("This jointstype is not implemented.")
+
+        if translation:
+            x_translations = x[:, -1, :3]
+            x_rotations = x[:, :-1]
+        else:
+            x_rotations = x
+
+        x_rotations = x_rotations.permute(0, 3, 1, 2)
+        nsamples, time, njoints, feats = x_rotations.shape
+
+        # Compute rotations (convert only masked sequences output)
+        if pose_rep == "rotvec":
+            rotations = geometry.axis_angle_to_matrix(x_rotations[mask])
+        elif pose_rep == "rotmat":
+            rotations = x_rotations[mask].view(-1, njoints, 3, 3)
+        elif pose_rep == "rotquat":
+            rotations = geometry.quaternion_to_matrix(x_rotations[mask])
+        elif pose_rep == "rot6d":
+            rotations = geometry.rotation_6d_to_matrix(x_rotations[mask])
+        else:
+            raise NotImplementedError("No geometry for this one.")
+
+        if not glob:
+            global_orient = torch.tensor(glob_rot, device=x.device)
+            global_orient = geometry.axis_angle_to_matrix(global_orient).view(1, 1, 3, 3)
+            global_orient = global_orient.repeat(len(rotations), 1, 1, 1)
+        else:
+            global_orient = rotations[:, 0]
+            rotations = rotations[:, 1:]
+
+        if betas is None:
+            betas = torch.zeros([rotations.shape[0], self.smpl_model.num_betas],
+                                dtype=rotations.dtype, device=rotations.device)
+            betas[:, 1] = beta
+            # import ipdb; ipdb.set_trace()
+        out = self.smpl_model(body_pose=rotations, global_orient=global_orient, betas=betas)
+
+        # get the desirable joints
+        joints = out[jointstype]
+
+        x_xyz = torch.empty(nsamples, time, joints.shape[1], 3, device=x.device, dtype=x.dtype)
+        x_xyz[~mask] = 0
+        x_xyz[mask] = joints
+
+        x_xyz = x_xyz.permute(0, 2, 3, 1).contiguous()
+
+        # the first translation root at the origin on the prediction
+        if jointstype != "vertices":
+            rootindex = JOINTSTYPE_ROOT[jointstype]
+            x_xyz = x_xyz - x_xyz[:, [rootindex], :, :]
+
+        if translation and vertstrans:
+            # the first translation root at the origin
+            x_translations = x_translations - x_translations[:, :, [0]]
+
+            # add the translation to all the joints
+            x_xyz = x_xyz + x_translations[:, None, :, :]
+
+        if get_rotations_back:
+            return x_xyz, rotations, global_orient
+        else:
+            return x_xyz

models/smpl.py

@@ -0,0 +1,97 @@
+# This code is based on https://github.com/Mathux/ACTOR.git
+import numpy as np
+import torch
+
+import contextlib
+
+from smplx import SMPLLayer as _SMPLLayer
+from smplx.lbs import vertices2joints
+
+
+# action2motion_joints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 21, 24, 38]
+# change 0 and 8
+action2motion_joints = [8, 1, 2, 3, 4, 5, 6, 7, 0, 9, 10, 11, 12, 13, 14, 21, 24, 38]
+
+from utils.config import SMPL_MODEL_PATH, JOINT_REGRESSOR_TRAIN_EXTRA
+
+JOINTSTYPE_ROOT = {"a2m": 0, # action2motion
+                   "smpl": 0,
+                   "a2mpl": 0, # set(smpl, a2m)
+                   "vibe": 8}  # 0 is the 8 position: OP MidHip below
+
+JOINT_MAP = {
+    'OP Nose': 24, 'OP Neck': 12, 'OP RShoulder': 17,
+    'OP RElbow': 19, 'OP RWrist': 21, 'OP LShoulder': 16,
+    'OP LElbow': 18, 'OP LWrist': 20, 'OP MidHip': 0,
+    'OP RHip': 2, 'OP RKnee': 5, 'OP RAnkle': 8,
+    'OP LHip': 1, 'OP LKnee': 4, 'OP LAnkle': 7,
+    'OP REye': 25, 'OP LEye': 26, 'OP REar': 27,
+    'OP LEar': 28, 'OP LBigToe': 29, 'OP LSmallToe': 30,
+    'OP LHeel': 31, 'OP RBigToe': 32, 'OP RSmallToe': 33, 'OP RHeel': 34,
+    'Right Ankle': 8, 'Right Knee': 5, 'Right Hip': 45,
+    'Left Hip': 46, 'Left Knee': 4, 'Left Ankle': 7,
+    'Right Wrist': 21, 'Right Elbow': 19, 'Right Shoulder': 17,
+    'Left Shoulder': 16, 'Left Elbow': 18, 'Left Wrist': 20,
+    'Neck (LSP)': 47, 'Top of Head (LSP)': 48,
+    'Pelvis (MPII)': 49, 'Thorax (MPII)': 50,
+    'Spine (H36M)': 51, 'Jaw (H36M)': 52,
+    'Head (H36M)': 53, 'Nose': 24, 'Left Eye': 26,
+    'Right Eye': 25, 'Left Ear': 28, 'Right Ear': 27
+}
+
+JOINT_NAMES = [
+    'OP Nose', 'OP Neck', 'OP RShoulder',
+    'OP RElbow', 'OP RWrist', 'OP LShoulder',
+    'OP LElbow', 'OP LWrist', 'OP MidHip',
+    'OP RHip', 'OP RKnee', 'OP RAnkle',
+    'OP LHip', 'OP LKnee', 'OP LAnkle',
+    'OP REye', 'OP LEye', 'OP REar',
+    'OP LEar', 'OP LBigToe', 'OP LSmallToe',
+    'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel',
+    'Right Ankle', 'Right Knee', 'Right Hip',
+    'Left Hip', 'Left Knee', 'Left Ankle',
+    'Right Wrist', 'Right Elbow', 'Right Shoulder',
+    'Left Shoulder', 'Left Elbow', 'Left Wrist',
+    'Neck (LSP)', 'Top of Head (LSP)',
+    'Pelvis (MPII)', 'Thorax (MPII)',
+    'Spine (H36M)', 'Jaw (H36M)',
+    'Head (H36M)', 'Nose', 'Left Eye',
+    'Right Eye', 'Left Ear', 'Right Ear'
+]
+
+
+# adapted from VIBE/SPIN to output smpl_joints, vibe joints and action2motion joints
+class SMPL(_SMPLLayer):
+    """ Extension of the official SMPL implementation to support more joints """
+
+    def __init__(self, model_path=SMPL_MODEL_PATH, **kwargs):
+        kwargs["model_path"] = model_path
+
+        # remove the verbosity for the 10-shapes beta parameters
+        with contextlib.redirect_stdout(None):
+            super(SMPL, self).__init__(**kwargs)
+            
+        J_regressor_extra = np.load(JOINT_REGRESSOR_TRAIN_EXTRA)
+        self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32))
+        vibe_indexes = np.array([JOINT_MAP[i] for i in JOINT_NAMES])
+        a2m_indexes = vibe_indexes[action2motion_joints]
+        smpl_indexes = np.arange(24)
+        a2mpl_indexes = np.unique(np.r_[smpl_indexes, a2m_indexes])
+
+        self.maps = {"vibe": vibe_indexes,
+                     "a2m": a2m_indexes,
+                     "smpl": smpl_indexes,
+                     "a2mpl": a2mpl_indexes}
+        
+    def forward(self, *args, **kwargs):
+        smpl_output = super(SMPL, self).forward(*args, **kwargs)
+        
+        extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)
+        all_joints = torch.cat([smpl_output.joints, extra_joints], dim=1)
+
+        output = {"vertices": smpl_output.vertices}
+
+        for joinstype, indexes in self.maps.items():
+            output[joinstype] = all_joints[:, indexes]
+            
+        return output

render_final.py

@@ -0,0 +1,189 @@
+from models.rotation2xyz import Rotation2xyz
+import numpy as np
+from trimesh import Trimesh
+import os
+os.environ['PYOPENGL_PLATFORM'] = "osmesa"
+
+import torch
+from visualize.simplify_loc2rot import joints2smpl
+import pyrender
+import matplotlib.pyplot as plt
+
+import io
+import imageio
+from shapely import geometry
+import trimesh
+from pyrender.constants import RenderFlags
+import math
+# import ffmpeg
+from PIL import Image
+import argparse
+
+class WeakPerspectiveCamera(pyrender.Camera):
+    def __init__(self,
+                 scale,
+                 translation,
+                 znear=pyrender.camera.DEFAULT_Z_NEAR,
+                 zfar=None,
+                 name=None):
+        super(WeakPerspectiveCamera, self).__init__(
+            znear=znear,
+            zfar=zfar,
+            name=name,
+        )
+        self.scale = scale
+        self.translation = translation
+
+    def get_projection_matrix(self, width=None, height=None):
+        P = np.eye(4)
+        P[0, 0] = self.scale[0]
+        P[1, 1] = self.scale[1]
+        P[0, 3] = self.translation[0] * self.scale[0]
+        P[1, 3] = -self.translation[1] * self.scale[1]
+        P[2, 2] = -1
+        return P
+
+def render(motions, outdir='test_vis', device_id=0, name=None, pred=True):
+    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
+
+    j2s = joints2smpl(num_frames=frames, device_id=device_id, cuda=False)
+    rot2xyz = Rotation2xyz(device=torch.device('cpu'))
+    faces = rot2xyz.smpl_model.faces
+
+    filepath = os.path.join(outdir, f'{name}.pt')
+    if not os.path.exists(filepath):
+        print(f'Running SMPLify, it may take a few minutes.')
+        motion_tensor, opt_dict = j2s.joint2smpl(motions)
+        vertices = rot2xyz(torch.tensor(motion_tensor).clone(), mask=None,
+                           pose_rep='rot6d', translation=True, glob=True,
+                           jointstype='vertices', vertstrans=True)
+        torch.save(vertices, filepath)
+    else:
+        vertices = torch.load(filepath)
+    
+    if not os.path.exists(outdir):
+        os.makedirs(outdir)
+    
+    frames = vertices.shape[3] # shape: 1, nb_frames, 3, nb_joints
+    print (vertices.shape)
+    MINS = torch.min(torch.min(vertices[0], axis=0)[0], axis=1)[0]
+    MAXS = torch.max(torch.max(vertices[0], axis=0)[0], axis=1)[0]
+    # vertices[:,:,1,:] -= MINS[1] + 1e-5
+
+
+    out_list = []
+    
+    minx = MINS[0] - 0.5
+    maxx = MAXS[0] + 0.5
+    minz = MINS[2] - 0.5 
+    maxz = MAXS[2] + 0.5
+    polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]])
+    polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5)
+
+    vid = []
+    for i in range(frames):
+        if i % 10 == 0:
+            print(i)
+
+        mesh = Trimesh(vertices=vertices[0, :, :, i].squeeze().tolist(), faces=faces)
+
+        base_color = (0.11, 0.53, 0.8, 0.5)
+        ## OPAQUE rendering without alpha
+        ## BLEND rendering consider alpha 
+        material = pyrender.MetallicRoughnessMaterial(
+            metallicFactor=0.7,
+            alphaMode='OPAQUE',
+            baseColorFactor=base_color
+        )
+
+
+        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
+
+        polygon_mesh.visual.face_colors = [0, 0, 0, 0.21]
+        polygon_render = pyrender.Mesh.from_trimesh(polygon_mesh, smooth=False)
+
+        bg_color = [1, 1, 1, 0.8]
+        scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4))
+        
+        sx, sy, tx, ty = [0.75, 0.75, 0, 0.10]
+
+        camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0))
+
+        light = pyrender.DirectionalLight(color=[1,1,1], intensity=300)
+
+        scene.add(mesh)
+
+        c = np.pi / 2
+
+        scene.add(polygon_render, pose=np.array([[ 1, 0, 0, 0],
+
+        [ 0, np.cos(c), -np.sin(c), MINS[1].cpu().numpy()],
+
+        [ 0, np.sin(c), np.cos(c), 0],
+
+        [ 0, 0, 0, 1]]))
+
+        light_pose = np.eye(4)
+        light_pose[:3, 3] = [0, -1, 1]
+        scene.add(light, pose=light_pose.copy())
+
+        light_pose[:3, 3] = [0, 1, 1]
+        scene.add(light, pose=light_pose.copy())
+
+        light_pose[:3, 3] = [1, 1, 2]
+        scene.add(light, pose=light_pose.copy())
+
+
+        c = -np.pi / 6
+
+        scene.add(camera, pose=[[ 1, 0, 0, (minx+maxx).cpu().numpy()/2],
+
+                                [ 0, np.cos(c), -np.sin(c), 1.5],
+
+                                [ 0, np.sin(c), np.cos(c), max(4, minz.cpu().numpy()+(1.5-MINS[1].cpu().numpy())*2, (maxx-minx).cpu().numpy())],
+
+                                [ 0, 0, 0, 1]
+                                ])
+        
+        # render scene
+        r = pyrender.OffscreenRenderer(960, 960)
+
+        color, _ = r.render(scene, flags=RenderFlags.RGBA)
+        # Image.fromarray(color).save(outdir+'/'+name+'_'+str(i)+'.png')
+
+        vid.append(color)
+
+        r.delete()
+
+    out = np.stack(vid, axis=0)
+    gif_path = os.path.join(outdir, f'{name}.gif')
+    imageio.mimsave(gif_path, out, fps=20)
+    
+
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Render 3D motion from a numpy file.")
+    parser.add_argument("filepath", type=str, help="Path to the numpy file containing the motion data.")
+    parser.add_argument("--outdir", type=str, default="./output", help="Output directory for rendered files.")
+    args = parser.parse_args()
+
+    # Ensure output directory exists
+    if not os.path.exists(args.outdir):
+        os.makedirs(args.outdir)
+
+    # Load the motion data from the provided file path
+    motions = np.load(args.filepath)
+    print('Loaded motion data from:', args.filepath, 'Shape:', motions.shape)
+
+    # Extract filename for naming output files
+    filename = os.path.basename(args.filepath).replace('.npy', '')
+
+    # Render the motion
+    render(motions[0], outdir=args.outdir, device_id=0, name=filename)

utils/config.py

@@ -0,0 +1,17 @@
+import os
+
+SMPL_DATA_PATH = "./body_models/smpl"
+
+SMPL_KINTREE_PATH = os.path.join(SMPL_DATA_PATH, "kintree_table.pkl")
+SMPL_MODEL_PATH = os.path.join(SMPL_DATA_PATH, "SMPL_NEUTRAL.pkl")
+JOINT_REGRESSOR_TRAIN_EXTRA = os.path.join(SMPL_DATA_PATH, 'J_regressor_extra.npy')
+
+ROT_CONVENTION_TO_ROT_NUMBER = {
+    'legacy': 23,
+    'no_hands': 21,
+    'full_hands': 51,
+    'mitten_hands': 33,
+}
+
+GENDERS = ['neutral', 'male', 'female']
+NUM_BETAS = 10

utils/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
+    
+    

visualization/plot_3d_global.py

@@ -0,0 +1,129 @@
+import torch 
+import matplotlib.pyplot as plt
+import numpy as np
+import io
+import matplotlib
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import mpl_toolkits.mplot3d.axes3d as p3
+from textwrap import wrap
+import imageio
+
+def plot_3d_motion(args, figsize=(10, 10), fps=120, radius=4):
+    matplotlib.use('Agg')
+    
+    
+    joints, out_name, title = args
+    
+    data = joints.copy().reshape(len(joints), -1, 3)
+    
+    nb_joints = joints.shape[1]
+    smpl_kinetic_chain = [[0, 11, 12, 13, 14, 15], [0, 16, 17, 18, 19, 20], [0, 1, 2, 3, 4], [3, 5, 6, 7], [3, 8, 9, 10]] if nb_joints == 21 else [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20]]
+    limits = 1000 if nb_joints == 21 else 2
+    MINS = data.min(axis=0).min(axis=0)
+    MAXS = data.max(axis=0).max(axis=0)
+    colors = ['red', 'blue', 'black', 'red', 'blue',
+              'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
+              'darkred', 'darkred', 'darkred', 'darkred', 'darkred']
+    frame_number = data.shape[0]
+    #     print(data.shape)
+
+    height_offset = MINS[1]
+    data[:, :, 1] -= height_offset
+    trajec = data[:, 0, [0, 2]]
+
+    data[..., 0] -= data[:, 0:1, 0]
+    data[..., 2] -= data[:, 0:1, 2]
+
+    def update(index):
+
+        def init():
+            ax.set_xlim(-limits, limits)
+            ax.set_ylim(-limits, limits)
+            ax.set_zlim(0, limits)
+            ax.grid(b=False)
+        def plot_xzPlane(minx, maxx, miny, minz, maxz):
+            ## Plot a plane XZ
+            verts = [
+                [minx, miny, minz],
+                [minx, miny, maxz],
+                [maxx, miny, maxz],
+                [maxx, miny, minz]
+            ]
+            xz_plane = Poly3DCollection([verts])
+            xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
+            ax.add_collection3d(xz_plane)
+        fig = plt.figure(figsize=(480/96., 320/96.), dpi=96) if nb_joints == 21 else plt.figure(figsize=(10, 10), dpi=96)
+        if title is not None :
+            wraped_title = '\n'.join(wrap(title, 40))
+            fig.suptitle(wraped_title, fontsize=16)
+        ax = p3.Axes3D(fig)
+        
+        init()
+        
+        ax.lines = []
+        ax.collections = []
+        ax.view_init(elev=110, azim=-90)
+        ax.dist = 7.5
+        #         ax =
+        plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, MINS[2] - trajec[index, 1],
+                     MAXS[2] - trajec[index, 1])
+        #         ax.scatter(data[index, :22, 0], data[index, :22, 1], data[index, :22, 2], color='black', s=3)
+
+        if index > 1:
+            ax.plot3D(trajec[:index, 0] - trajec[index, 0], np.zeros_like(trajec[:index, 0]),
+                      trajec[:index, 1] - trajec[index, 1], linewidth=1.0,
+                      color='blue')
+        #             ax = plot_xzPlane(ax, MINS[0], MAXS[0], 0, MINS[2], MAXS[2])
+
+        for i, (chain, color) in enumerate(zip(smpl_kinetic_chain, colors)):
+            #             print(color)
+            if i < 5:
+                linewidth = 4.0
+            else:
+                linewidth = 2.0
+            ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth,
+                      color=color)
+        #         print(trajec[:index, 0].shape)
+
+        plt.axis('off')
+        ax.set_xticklabels([])
+        ax.set_yticklabels([])
+        ax.set_zticklabels([])
+    
+        if out_name is not None : 
+            plt.savefig(out_name, dpi=96)
+            plt.close()
+            
+        else : 
+            io_buf = io.BytesIO()
+            fig.savefig(io_buf, format='raw', dpi=96)
+            io_buf.seek(0)
+            # print(fig.bbox.bounds)
+            arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8),
+                                newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1))
+            io_buf.close()
+            plt.close()
+            return arr
+
+    out = []
+    for i in range(frame_number) : 
+        out.append(update(i))
+    out = np.stack(out, axis=0)
+    return torch.from_numpy(out)
+
+
+def draw_to_batch(smpl_joints_batch, title_batch=None, outname=None) : 
+    
+    batch_size = len(smpl_joints_batch)
+    out = []
+    for i in range(batch_size) : 
+        out.append(plot_3d_motion([smpl_joints_batch[i], None, title_batch[i] if title_batch is not None else None]))
+        if outname is not None:
+            imageio.mimsave(outname[i], np.array(out[-1]), fps=20)
+    out = torch.stack(out, axis=0)
+    return out
+    
+
+
+
+

visualize/joints2smpl/smpl_models/SMPL_downsample_index.pkl

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

visualize/joints2smpl/smpl_models/gmm_08.pkl

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

visualize/joints2smpl/smpl_models/neutral_smpl_mean_params.h5

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

visualize/joints2smpl/smpl_models/smplx_parts_segm.pkl

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

visualize/joints2smpl/src/config.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+# Map joints Name to SMPL joints idx
+JOINT_MAP = {
+'MidHip': 0,
+'LHip': 1, 'LKnee': 4, 'LAnkle': 7, 'LFoot': 10,
+'RHip': 2, 'RKnee': 5, 'RAnkle': 8, 'RFoot': 11,
+'LShoulder': 16, 'LElbow': 18, 'LWrist': 20, 'LHand': 22, 
+'RShoulder': 17, 'RElbow': 19, 'RWrist': 21, 'RHand': 23,
+'spine1': 3, 'spine2': 6, 'spine3': 9,  'Neck': 12, 'Head': 15,
+'LCollar':13, 'Rcollar' :14, 
+'Nose':24, 'REye':26,  'LEye':26,  'REar':27,  'LEar':28, 
+'LHeel': 31, 'RHeel': 34,
+'OP RShoulder': 17, 'OP LShoulder': 16,
+'OP RHip': 2, 'OP LHip': 1,
+'OP Neck': 12,
+}
+
+full_smpl_idx = range(24)
+key_smpl_idx = [0, 1, 4, 7,  2, 5, 8,  17, 19, 21,  16, 18, 20]
+
+
+AMASS_JOINT_MAP = {
+'MidHip': 0,
+'LHip': 1, 'LKnee': 4, 'LAnkle': 7, 'LFoot': 10,
+'RHip': 2, 'RKnee': 5, 'RAnkle': 8, 'RFoot': 11,
+'LShoulder': 16, 'LElbow': 18, 'LWrist': 20,  
+'RShoulder': 17, 'RElbow': 19, 'RWrist': 21, 
+'spine1': 3, 'spine2': 6, 'spine3': 9,  'Neck': 12, 'Head': 15,
+'LCollar':13, 'Rcollar' :14, 
+}
+amass_idx =       range(22)
+amass_smpl_idx =  range(22)
+
+
+SMPL_MODEL_DIR = "./body_models/"
+GMM_MODEL_DIR = "./visualize/joints2smpl/smpl_models/"
+SMPL_MEAN_FILE = "./visualize/joints2smpl/smpl_models/neutral_smpl_mean_params.h5"
+# for collsion 
+Part_Seg_DIR = "./visualize/joints2smpl/smpl_models/smplx_parts_segm.pkl"

visualize/joints2smpl/src/customloss.py

@@ -0,0 +1,222 @@
+import torch
+import torch.nn.functional as F
+from visualize.joints2smpl.src import config
+
+# Guassian
+def gmof(x, sigma):
+    """
+    Geman-McClure error function
+    """
+    x_squared = x ** 2
+    sigma_squared = sigma ** 2
+    return (sigma_squared * x_squared) / (sigma_squared + x_squared)
+
+# angle prior
+def angle_prior(pose):
+    """
+    Angle prior that penalizes unnatural bending of the knees and elbows
+    """
+    # We subtract 3 because pose does not include the global rotation of the model
+    return torch.exp(
+        pose[:, [55 - 3, 58 - 3, 12 - 3, 15 - 3]] * torch.tensor([1., -1., -1, -1.], device=pose.device)) ** 2
+
+
+def perspective_projection(points, rotation, translation,
+                           focal_length, camera_center):
+    """
+    This function computes the perspective projection of a set of points.
+    Input:
+        points (bs, N, 3): 3D points
+        rotation (bs, 3, 3): Camera rotation
+        translation (bs, 3): Camera translation
+        focal_length (bs,) or scalar: Focal length
+        camera_center (bs, 2): Camera center
+    """
+    batch_size = points.shape[0]
+    K = torch.zeros([batch_size, 3, 3], device=points.device)
+    K[:, 0, 0] = focal_length
+    K[:, 1, 1] = focal_length
+    K[:, 2, 2] = 1.
+    K[:, :-1, -1] = camera_center
+
+    # Transform points
+    points = torch.einsum('bij,bkj->bki', rotation, points)
+    points = points + translation.unsqueeze(1)
+
+    # Apply perspective distortion
+    projected_points = points / points[:, :, -1].unsqueeze(-1)
+
+    # Apply camera intrinsics
+    projected_points = torch.einsum('bij,bkj->bki', K, projected_points)
+
+    return projected_points[:, :, :-1]
+
+
+def body_fitting_loss(body_pose, betas, model_joints, camera_t, camera_center,
+                      joints_2d, joints_conf, pose_prior,
+                      focal_length=5000, sigma=100, pose_prior_weight=4.78,
+                      shape_prior_weight=5, angle_prior_weight=15.2,
+                      output='sum'):
+    """
+    Loss function for body fitting
+    """
+    batch_size = body_pose.shape[0]
+    rotation = torch.eye(3, device=body_pose.device).unsqueeze(0).expand(batch_size, -1, -1)
+
+    projected_joints = perspective_projection(model_joints, rotation, camera_t,
+                                              focal_length, camera_center)
+
+    # Weighted robust reprojection error
+    reprojection_error = gmof(projected_joints - joints_2d, sigma)
+    reprojection_loss = (joints_conf ** 2) * reprojection_error.sum(dim=-1)
+
+    # Pose prior loss
+    pose_prior_loss = (pose_prior_weight ** 2) * pose_prior(body_pose, betas)
+
+    # Angle prior for knees and elbows
+    angle_prior_loss = (angle_prior_weight ** 2) * angle_prior(body_pose).sum(dim=-1)
+
+    # Regularizer to prevent betas from taking large values
+    shape_prior_loss = (shape_prior_weight ** 2) * (betas ** 2).sum(dim=-1)
+
+    total_loss = reprojection_loss.sum(dim=-1) + pose_prior_loss + angle_prior_loss + shape_prior_loss
+
+    if output == 'sum':
+        return total_loss.sum()
+    elif output == 'reprojection':
+        return reprojection_loss
+
+
+# --- get camera fitting loss -----
+def camera_fitting_loss(model_joints, camera_t, camera_t_est, camera_center, 
+                        joints_2d, joints_conf,
+                        focal_length=5000, depth_loss_weight=100):
+    """
+    Loss function for camera optimization.
+    """
+    # Project model joints
+    batch_size = model_joints.shape[0]
+    rotation = torch.eye(3, device=model_joints.device).unsqueeze(0).expand(batch_size, -1, -1)
+    projected_joints = perspective_projection(model_joints, rotation, camera_t,
+                                              focal_length, camera_center)
+
+    # get the indexed four
+    op_joints = ['OP RHip', 'OP LHip', 'OP RShoulder', 'OP LShoulder']
+    op_joints_ind = [config.JOINT_MAP[joint] for joint in op_joints]
+    gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder']
+    gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+
+    reprojection_error_op = (joints_2d[:, op_joints_ind] -
+                             projected_joints[:, op_joints_ind]) ** 2
+    reprojection_error_gt = (joints_2d[:, gt_joints_ind] -
+                             projected_joints[:, gt_joints_ind]) ** 2
+
+    # Check if for each example in the batch all 4 OpenPose detections are valid, otherwise use the GT detections
+    # OpenPose joints are more reliable for this task, so we prefer to use them if possible
+    is_valid = (joints_conf[:, op_joints_ind].min(dim=-1)[0][:, None, None] > 0).float()
+    reprojection_loss = (is_valid * reprojection_error_op + (1 - is_valid) * reprojection_error_gt).sum(dim=(1, 2))
+
+    # Loss that penalizes deviation from depth estimate
+    depth_loss = (depth_loss_weight ** 2) * (camera_t[:, 2] - camera_t_est[:, 2]) ** 2
+
+    total_loss = reprojection_loss + depth_loss
+    return total_loss.sum()
+
+
+
+ # #####--- body fitiing loss -----
+def body_fitting_loss_3d(body_pose, preserve_pose,
+                         betas, model_joints, camera_translation,
+                         j3d, pose_prior,
+                         joints3d_conf,
+                         sigma=100, pose_prior_weight=4.78*1.5,
+                         shape_prior_weight=5.0, angle_prior_weight=15.2,
+                         joint_loss_weight=500.0,
+                         pose_preserve_weight=0.0,
+                         use_collision=False,
+                         model_vertices=None, model_faces=None,
+                         search_tree=None,  pen_distance=None,  filter_faces=None,
+                         collision_loss_weight=1000
+                         ):
+    """
+    Loss function for body fitting
+    """
+    batch_size = body_pose.shape[0]
+
+    #joint3d_loss = (joint_loss_weight ** 2) * gmof((model_joints + camera_translation) - j3d, sigma).sum(dim=-1)
+    
+    joint3d_error = gmof((model_joints + camera_translation) - j3d, sigma)
+    
+    joint3d_loss_part = (joints3d_conf ** 2) * joint3d_error.sum(dim=-1)
+    joint3d_loss = ((joint_loss_weight ** 2) * joint3d_loss_part).sum(dim=-1)
+    
+    # Pose prior loss
+    pose_prior_loss = (pose_prior_weight ** 2) * pose_prior(body_pose, betas)
+    # Angle prior for knees and elbows
+    angle_prior_loss = (angle_prior_weight ** 2) * angle_prior(body_pose).sum(dim=-1)
+    # Regularizer to prevent betas from taking large values
+    shape_prior_loss = (shape_prior_weight ** 2) * (betas ** 2).sum(dim=-1)
+
+    collision_loss = 0.0
+    # Calculate the loss due to interpenetration
+    if use_collision:
+        triangles = torch.index_select(
+            model_vertices, 1,
+            model_faces).view(batch_size, -1, 3, 3)
+
+        with torch.no_grad():
+            collision_idxs = search_tree(triangles)
+
+        # Remove unwanted collisions
+        if filter_faces is not None:
+            collision_idxs = filter_faces(collision_idxs)
+
+        if collision_idxs.ge(0).sum().item() > 0:
+            collision_loss = torch.sum(collision_loss_weight * pen_distance(triangles, collision_idxs))
+    
+    pose_preserve_loss = (pose_preserve_weight ** 2) * ((body_pose - preserve_pose) ** 2).sum(dim=-1)
+
+    # print('joint3d_loss', joint3d_loss.shape)
+    # print('pose_prior_loss', pose_prior_loss.shape)
+    # print('angle_prior_loss', angle_prior_loss.shape)
+    # print('shape_prior_loss', shape_prior_loss.shape)
+    # print('collision_loss', collision_loss)
+    # print('pose_preserve_loss', pose_preserve_loss.shape)
+
+    total_loss = joint3d_loss + pose_prior_loss + angle_prior_loss + shape_prior_loss + collision_loss + pose_preserve_loss
+
+    return total_loss.sum()
+
+
+# #####--- get camera fitting loss -----
+def camera_fitting_loss_3d(model_joints, camera_t, camera_t_est,
+                           j3d, joints_category="orig", depth_loss_weight=100.0):
+    """
+    Loss function for camera optimization.
+    """
+    model_joints = model_joints + camera_t
+    # # get the indexed four
+    # op_joints = ['OP RHip', 'OP LHip', 'OP RShoulder', 'OP LShoulder']
+    # op_joints_ind = [config.JOINT_MAP[joint] for joint in op_joints]
+    #
+    # j3d_error_loss = (j3d[:, op_joints_ind] -
+    #                          model_joints[:, op_joints_ind]) ** 2
+
+    gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder']
+    gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+    
+    if joints_category=="orig":
+        select_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+    elif joints_category=="AMASS":
+        select_joints_ind = [config.AMASS_JOINT_MAP[joint] for joint in gt_joints]
+    else:
+        print("NO SUCH JOINTS CATEGORY!")
+
+    j3d_error_loss = (j3d[:, select_joints_ind] -
+                      model_joints[:, gt_joints_ind]) ** 2
+
+    # Loss that penalizes deviation from depth estimate
+    depth_loss = (depth_loss_weight**2) *  (camera_t - camera_t_est)**2
+
+    total_loss = j3d_error_loss +  depth_loss
+    return total_loss.sum()

visualize/joints2smpl/src/prior.py

@@ -0,0 +1,230 @@
+# -*- 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
+
+import sys
+import os
+
+import time
+import pickle
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+DEFAULT_DTYPE = torch.float32
+
+
+def create_prior(prior_type, **kwargs):
+    if prior_type == 'gmm':
+        prior = MaxMixturePrior(**kwargs)
+    elif prior_type == 'l2':
+        return L2Prior(**kwargs)
+    elif prior_type == 'angle':
+        return SMPLifyAnglePrior(**kwargs)
+    elif prior_type == 'none' or prior_type is None:
+        # Don't use any pose prior
+        def no_prior(*args, **kwargs):
+            return 0.0
+        prior = no_prior
+    else:
+        raise ValueError('Prior {}'.format(prior_type) + ' is not implemented')
+    return prior
+
+
+class SMPLifyAnglePrior(nn.Module):
+    def __init__(self, dtype=torch.float32, **kwargs):
+        super(SMPLifyAnglePrior, self).__init__()
+
+        # Indices for the roration angle of
+        # 55: left elbow,  90deg bend at -np.pi/2
+        # 58: right elbow, 90deg bend at np.pi/2
+        # 12: left knee,   90deg bend at np.pi/2
+        # 15: right knee,  90deg bend at np.pi/2
+        angle_prior_idxs = np.array([55, 58, 12, 15], dtype=np.int64)
+        angle_prior_idxs = torch.tensor(angle_prior_idxs, dtype=torch.long)
+        self.register_buffer('angle_prior_idxs', angle_prior_idxs)
+
+        angle_prior_signs = np.array([1, -1, -1, -1],
+                                     dtype=np.float32 if dtype == torch.float32
+                                     else np.float64)
+        angle_prior_signs = torch.tensor(angle_prior_signs,
+                                         dtype=dtype)
+        self.register_buffer('angle_prior_signs', angle_prior_signs)
+
+    def forward(self, pose, with_global_pose=False):
+        ''' Returns the angle prior loss for the given pose
+
+        Args:
+            pose: (Bx[23 + 1] * 3) torch tensor with the axis-angle
+            representation of the rotations of the joints of the SMPL model.
+        Kwargs:
+            with_global_pose: Whether the pose vector also contains the global
+            orientation of the SMPL model. If not then the indices must be
+            corrected.
+        Returns:
+            A sze (B) tensor containing the angle prior loss for each element
+            in the batch.
+        '''
+        angle_prior_idxs = self.angle_prior_idxs - (not with_global_pose) * 3
+        return torch.exp(pose[:, angle_prior_idxs] *
+                         self.angle_prior_signs).pow(2)
+
+
+class L2Prior(nn.Module):
+    def __init__(self, dtype=DEFAULT_DTYPE, reduction='sum', **kwargs):
+        super(L2Prior, self).__init__()
+
+    def forward(self, module_input, *args):
+        return torch.sum(module_input.pow(2))
+
+
+class MaxMixturePrior(nn.Module):
+
+    def __init__(self, prior_folder='prior',
+                 num_gaussians=6, dtype=DEFAULT_DTYPE, epsilon=1e-16,
+                 use_merged=True,
+                 **kwargs):
+        super(MaxMixturePrior, self).__init__()
+
+        if dtype == DEFAULT_DTYPE:
+            np_dtype = np.float32
+        elif dtype == torch.float64:
+            np_dtype = np.float64
+        else:
+            print('Unknown float type {}, exiting!'.format(dtype))
+            sys.exit(-1)
+
+        self.num_gaussians = num_gaussians
+        self.epsilon = epsilon
+        self.use_merged = use_merged
+        gmm_fn = 'gmm_{:02d}.pkl'.format(num_gaussians)
+
+        full_gmm_fn = os.path.join(prior_folder, gmm_fn)
+        if not os.path.exists(full_gmm_fn):
+            print('The path to the mixture prior "{}"'.format(full_gmm_fn) +
+                  ' does not exist, exiting!')
+            sys.exit(-1)
+
+        with open(full_gmm_fn, 'rb') as f:
+            gmm = pickle.load(f, encoding='latin1')
+
+        if type(gmm) == dict:
+            means = gmm['means'].astype(np_dtype)
+            covs = gmm['covars'].astype(np_dtype)
+            weights = gmm['weights'].astype(np_dtype)
+        elif 'sklearn.mixture.gmm.GMM' in str(type(gmm)):
+            means = gmm.means_.astype(np_dtype)
+            covs = gmm.covars_.astype(np_dtype)
+            weights = gmm.weights_.astype(np_dtype)
+        else:
+            print('Unknown type for the prior: {}, exiting!'.format(type(gmm)))
+            sys.exit(-1)
+
+        self.register_buffer('means', torch.tensor(means, dtype=dtype))
+
+        self.register_buffer('covs', torch.tensor(covs, dtype=dtype))
+
+        precisions = [np.linalg.inv(cov) for cov in covs]
+        precisions = np.stack(precisions).astype(np_dtype)
+
+        self.register_buffer('precisions',
+                             torch.tensor(precisions, dtype=dtype))
+
+        # The constant term:
+        sqrdets = np.array([(np.sqrt(np.linalg.det(c)))
+                            for c in gmm['covars']])
+        const = (2 * np.pi)**(69 / 2.)
+
+        nll_weights = np.asarray(gmm['weights'] / (const *
+                                                   (sqrdets / sqrdets.min())))
+        nll_weights = torch.tensor(nll_weights, dtype=dtype).unsqueeze(dim=0)
+        self.register_buffer('nll_weights', nll_weights)
+
+        weights = torch.tensor(gmm['weights'], dtype=dtype).unsqueeze(dim=0)
+        self.register_buffer('weights', weights)
+
+        self.register_buffer('pi_term',
+                             torch.log(torch.tensor(2 * np.pi, dtype=dtype)))
+
+        cov_dets = [np.log(np.linalg.det(cov.astype(np_dtype)) + epsilon)
+                    for cov in covs]
+        self.register_buffer('cov_dets',
+                             torch.tensor(cov_dets, dtype=dtype))
+
+        # The dimensionality of the random variable
+        self.random_var_dim = self.means.shape[1]
+
+    def get_mean(self):
+        ''' Returns the mean of the mixture '''
+        mean_pose = torch.matmul(self.weights, self.means)
+        return mean_pose
+
+    def merged_log_likelihood(self, pose, betas):
+        diff_from_mean = pose.unsqueeze(dim=1) - self.means
+
+        prec_diff_prod = torch.einsum('mij,bmj->bmi',
+                                      [self.precisions, diff_from_mean])
+        diff_prec_quadratic = (prec_diff_prod * diff_from_mean).sum(dim=-1)
+
+        curr_loglikelihood = 0.5 * diff_prec_quadratic - \
+            torch.log(self.nll_weights)
+        #  curr_loglikelihood = 0.5 * (self.cov_dets.unsqueeze(dim=0) +
+        #  self.random_var_dim * self.pi_term +
+        #  diff_prec_quadratic
+        #  ) - torch.log(self.weights)
+
+        min_likelihood, _ = torch.min(curr_loglikelihood, dim=1)
+        return min_likelihood
+
+    def log_likelihood(self, pose, betas, *args, **kwargs):
+        ''' Create graph operation for negative log-likelihood calculation
+        '''
+        likelihoods = []
+
+        for idx in range(self.num_gaussians):
+            mean = self.means[idx]
+            prec = self.precisions[idx]
+            cov = self.covs[idx]
+            diff_from_mean = pose - mean
+
+            curr_loglikelihood = torch.einsum('bj,ji->bi',
+                                              [diff_from_mean, prec])
+            curr_loglikelihood = torch.einsum('bi,bi->b',
+                                              [curr_loglikelihood,
+                                               diff_from_mean])
+            cov_term = torch.log(torch.det(cov) + self.epsilon)
+            curr_loglikelihood += 0.5 * (cov_term +
+                                         self.random_var_dim *
+                                         self.pi_term)
+            likelihoods.append(curr_loglikelihood)
+
+        log_likelihoods = torch.stack(likelihoods, dim=1)
+        min_idx = torch.argmin(log_likelihoods, dim=1)
+        weight_component = self.nll_weights[:, min_idx]
+        weight_component = -torch.log(weight_component)
+
+        return weight_component + log_likelihoods[:, min_idx]
+
+    def forward(self, pose, betas):
+        if self.use_merged:
+            return self.merged_log_likelihood(pose, betas)
+        else:
+            return self.log_likelihood(pose, betas)

visualize/joints2smpl/src/smplify.py

@@ -0,0 +1,279 @@
+import torch
+import os, sys
+import pickle
+import smplx
+import numpy as np
+
+sys.path.append(os.path.dirname(__file__))
+from customloss import (camera_fitting_loss, 
+                        body_fitting_loss, 
+                        camera_fitting_loss_3d,
+                        body_fitting_loss_3d, 
+                        )
+from prior import MaxMixturePrior
+from visualize.joints2smpl.src import config
+
+
+
+@torch.no_grad()
+def guess_init_3d(model_joints, 
+                  j3d, 
+                  joints_category="orig"):
+    """Initialize the camera translation via triangle similarity, by using the torso joints        .
+    :param model_joints: SMPL model with pre joints
+    :param j3d: 25x3 array of Kinect Joints
+    :returns: 3D vector corresponding to the estimated camera translation
+    """
+    # get the indexed four
+    gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder']
+    gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+    
+    if joints_category=="orig":
+        joints_ind_category = [config.JOINT_MAP[joint] for joint in gt_joints]
+    elif joints_category=="AMASS":
+        joints_ind_category = [config.AMASS_JOINT_MAP[joint] for joint in gt_joints] 
+    else:
+        print("NO SUCH JOINTS CATEGORY!") 
+
+    sum_init_t = (j3d[:, joints_ind_category] - model_joints[:, gt_joints_ind]).sum(dim=1)
+    init_t = sum_init_t / 4.0
+    return init_t
+
+
+# SMPLIfy 3D
+class SMPLify3D():
+    """Implementation of SMPLify, use 3D joints."""
+
+    def __init__(self,
+                 smplxmodel,
+                 step_size=1e-2,
+                 batch_size=1,
+                 num_iters=100,
+                 use_collision=False,
+                 use_lbfgs=True,
+                 joints_category="orig",
+                 device=torch.device('cuda:0'),
+                 ):
+
+        # Store options
+        self.batch_size = batch_size
+        self.device = device
+        self.step_size = step_size
+
+        self.num_iters = num_iters
+        # --- choose optimizer
+        self.use_lbfgs = use_lbfgs
+        # GMM pose prior
+        self.pose_prior = MaxMixturePrior(prior_folder=config.GMM_MODEL_DIR,
+                                          num_gaussians=8,
+                                          dtype=torch.float32).to(device)
+        # collision part
+        self.use_collision = use_collision
+        if self.use_collision:
+            self.part_segm_fn = config.Part_Seg_DIR
+        
+        # reLoad SMPL-X model
+        self.smpl = smplxmodel
+
+        self.model_faces = smplxmodel.faces_tensor.view(-1)
+
+        # select joint joint_category
+        self.joints_category = joints_category
+        
+        if joints_category=="orig":
+            self.smpl_index = config.full_smpl_idx
+            self.corr_index = config.full_smpl_idx 
+        elif joints_category=="AMASS":
+            self.smpl_index = config.amass_smpl_idx
+            self.corr_index = config.amass_idx
+        else:
+            self.smpl_index = None 
+            self.corr_index = None
+            print("NO SUCH JOINTS CATEGORY!")
+
+    # ---- get the man function here ------
+    def __call__(self, init_pose, init_betas, init_cam_t, j3d, conf_3d=1.0, seq_ind=0):
+        """Perform body fitting.
+        Input:
+            init_pose: SMPL pose estimate
+            init_betas: SMPL betas estimate
+            init_cam_t: Camera translation estimate
+            j3d: joints 3d aka keypoints
+            conf_3d: confidence for 3d joints
+			seq_ind: index of the sequence
+        Returns:
+            vertices: Vertices of optimized shape
+            joints: 3D joints of optimized shape
+            pose: SMPL pose parameters of optimized shape
+            betas: SMPL beta parameters of optimized shape
+            camera_translation: Camera translation
+        """
+
+        # # # add the mesh inter-section to avoid
+        search_tree = None
+        pen_distance = None
+        filter_faces = None
+        
+        if self.use_collision:
+            from mesh_intersection.bvh_search_tree import BVH
+            import mesh_intersection.loss as collisions_loss
+            from mesh_intersection.filter_faces import FilterFaces
+
+            search_tree = BVH(max_collisions=8)
+
+            pen_distance = collisions_loss.DistanceFieldPenetrationLoss(
+                           sigma=0.5, point2plane=False, vectorized=True, penalize_outside=True)
+
+            if self.part_segm_fn:
+                # Read the part segmentation
+                part_segm_fn = os.path.expandvars(self.part_segm_fn)
+                with open(part_segm_fn, 'rb') as faces_parents_file:
+                    face_segm_data = pickle.load(faces_parents_file,  encoding='latin1')
+                faces_segm = face_segm_data['segm']
+                faces_parents = face_segm_data['parents']
+                # Create the module used to filter invalid collision pairs
+                filter_faces = FilterFaces(
+                    faces_segm=faces_segm, faces_parents=faces_parents,
+                    ign_part_pairs=None).to(device=self.device)
+                    
+                    
+        # Split SMPL pose to body pose and global orientation
+        body_pose = init_pose[:, 3:].detach().clone()
+        global_orient = init_pose[:, :3].detach().clone()
+        betas = init_betas.detach().clone()
+
+        # use guess 3d to get the initial
+        smpl_output = self.smpl(global_orient=global_orient,
+                                body_pose=body_pose,
+                                betas=betas)
+        model_joints = smpl_output.joints
+
+        init_cam_t = guess_init_3d(model_joints, j3d, self.joints_category).unsqueeze(1).detach()
+        camera_translation = init_cam_t.clone()
+        
+        preserve_pose = init_pose[:, 3:].detach().clone()
+       # -------------Step 1: Optimize camera translation and body orientation--------
+        # Optimize only camera translation and body orientation
+        body_pose.requires_grad = False
+        betas.requires_grad = False
+        global_orient.requires_grad = True
+        camera_translation.requires_grad = True
+
+        camera_opt_params = [global_orient, camera_translation]
+
+        if self.use_lbfgs:
+            camera_optimizer = torch.optim.LBFGS(camera_opt_params, max_iter=self.num_iters,
+                                                 lr=self.step_size, line_search_fn='strong_wolfe')
+            for i in range(10):
+                def closure():
+                    camera_optimizer.zero_grad()
+                    smpl_output = self.smpl(global_orient=global_orient,
+                                            body_pose=body_pose,
+                                            betas=betas)
+                    model_joints = smpl_output.joints
+                    # print('model_joints', model_joints.shape)
+                    # print('camera_translation', camera_translation.shape)
+                    # print('init_cam_t', init_cam_t.shape)
+                    # print('j3d', j3d.shape)
+                    loss = camera_fitting_loss_3d(model_joints, camera_translation,
+                                                  init_cam_t, j3d, self.joints_category)
+                    loss.backward()
+                    return loss
+
+                camera_optimizer.step(closure)
+        else:
+            camera_optimizer = torch.optim.Adam(camera_opt_params, lr=self.step_size, betas=(0.9, 0.999))
+
+            for i in range(20):
+                smpl_output = self.smpl(global_orient=global_orient,
+                                        body_pose=body_pose,
+                                        betas=betas)
+                model_joints = smpl_output.joints
+
+                loss = camera_fitting_loss_3d(model_joints[:, self.smpl_index], camera_translation,
+                                              init_cam_t,  j3d[:, self.corr_index], self.joints_category)
+                camera_optimizer.zero_grad()
+                loss.backward()
+                camera_optimizer.step()
+
+        # Fix camera translation after optimizing camera
+        # --------Step 2: Optimize body joints --------------------------
+        # Optimize only the body pose and global orientation of the body
+        body_pose.requires_grad = True
+        global_orient.requires_grad = True
+        camera_translation.requires_grad = True
+
+        # --- if we use the sequence, fix the shape
+        if seq_ind == 0:
+            betas.requires_grad = True
+            body_opt_params = [body_pose, betas, global_orient, camera_translation]
+        else:
+            betas.requires_grad = False
+            body_opt_params = [body_pose, global_orient, camera_translation]
+
+        if self.use_lbfgs:
+            body_optimizer = torch.optim.LBFGS(body_opt_params, max_iter=self.num_iters,
+                                               lr=self.step_size, line_search_fn='strong_wolfe')
+            for i in range(self.num_iters):
+                def closure():
+                    body_optimizer.zero_grad()
+                    smpl_output = self.smpl(global_orient=global_orient,
+                                            body_pose=body_pose,
+                                            betas=betas)
+                    model_joints = smpl_output.joints
+                    model_vertices = smpl_output.vertices
+
+                    loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation,
+                                                j3d[:, self.corr_index], self.pose_prior,
+                                                joints3d_conf=conf_3d,
+                                                joint_loss_weight=600.0,
+                                                pose_preserve_weight=5.0,
+                                                use_collision=self.use_collision, 
+                                                model_vertices=model_vertices, model_faces=self.model_faces,
+                                                search_tree=search_tree, pen_distance=pen_distance, filter_faces=filter_faces)
+                    loss.backward()
+                    return loss
+
+                body_optimizer.step(closure)
+        else:
+            body_optimizer = torch.optim.Adam(body_opt_params, lr=self.step_size, betas=(0.9, 0.999))
+
+            for i in range(self.num_iters):
+                smpl_output = self.smpl(global_orient=global_orient,
+                                        body_pose=body_pose,
+                                        betas=betas)
+                model_joints = smpl_output.joints
+                model_vertices = smpl_output.vertices
+
+                loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation,
+                                            j3d[:, self.corr_index], self.pose_prior,
+                                            joints3d_conf=conf_3d,
+                                            joint_loss_weight=600.0,
+                                            use_collision=self.use_collision, 
+                                            model_vertices=model_vertices, model_faces=self.model_faces,
+                                            search_tree=search_tree,  pen_distance=pen_distance,  filter_faces=filter_faces)
+                body_optimizer.zero_grad()
+                loss.backward()
+                body_optimizer.step()
+
+        # Get final loss value
+        with torch.no_grad():
+            smpl_output = self.smpl(global_orient=global_orient,
+                                    body_pose=body_pose,
+                                    betas=betas, return_full_pose=True)
+            model_joints = smpl_output.joints
+            model_vertices = smpl_output.vertices
+
+            final_loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation,
+                                              j3d[:, self.corr_index], self.pose_prior,
+                                              joints3d_conf=conf_3d,
+                                              joint_loss_weight=600.0,
+                                              use_collision=self.use_collision, model_vertices=model_vertices, model_faces=self.model_faces,
+                                              search_tree=search_tree,  pen_distance=pen_distance,  filter_faces=filter_faces)
+
+        vertices = smpl_output.vertices.detach()
+        joints = smpl_output.joints.detach()
+        pose = torch.cat([global_orient, body_pose], dim=-1).detach()
+        betas = betas.detach()
+
+        return vertices, joints, pose, betas, camera_translation, final_loss

visualize/render_mesh.py

@@ -0,0 +1,33 @@
+import argparse
+import os
+from visualize import vis_utils
+import shutil
+from tqdm import tqdm
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_path", type=str, required=True, help='stick figure mp4 file to be rendered.')
+    parser.add_argument("--cuda", type=bool, default=True, help='')
+    parser.add_argument("--device", type=int, default=0, help='')
+    params = parser.parse_args()
+
+    assert params.input_path.endswith('.mp4')
+    parsed_name = os.path.basename(params.input_path).replace('.mp4', '').replace('sample', '').replace('rep', '')
+    sample_i, rep_i = [int(e) for e in parsed_name.split('_')]
+    npy_path = os.path.join(os.path.dirname(params.input_path), 'results.npy')
+    out_npy_path = params.input_path.replace('.mp4', '_smpl_params.npy')
+    assert os.path.exists(npy_path)
+    results_dir = params.input_path.replace('.mp4', '_obj')
+    if os.path.exists(results_dir):
+        shutil.rmtree(results_dir)
+    os.makedirs(results_dir)
+
+    npy2obj = vis_utils.npy2obj(npy_path, sample_i, rep_i,
+                                device=params.device, cuda=params.cuda)
+
+    print('Saving obj files to [{}]'.format(os.path.abspath(results_dir)))
+    for frame_i in tqdm(range(npy2obj.real_num_frames)):
+        npy2obj.save_obj(os.path.join(results_dir, 'frame{:03d}.obj'.format(frame_i)), frame_i)
+
+    print('Saving SMPL params to [{}]'.format(os.path.abspath(out_npy_path)))
+    npy2obj.save_npy(out_npy_path)

visualize/simplify_loc2rot.py

@@ -0,0 +1,131 @@
+import numpy as np
+import os
+import torch
+from visualize.joints2smpl.src import config
+import smplx
+import h5py
+from visualize.joints2smpl.src.smplify import SMPLify3D
+from tqdm import tqdm
+import utils.rotation_conversions as geometry
+import argparse
+
+
+class joints2smpl:
+
+    def __init__(self, num_frames, device_id, cuda=True):
+        self.device = torch.device("cuda:" + str(device_id) if cuda else "cpu")
+        # self.device = torch.device("cpu")
+        self.batch_size = num_frames
+        self.num_joints = 22  # for HumanML3D
+        self.joint_category = "AMASS"
+        self.num_smplify_iters = 120
+        self.fix_foot = False
+        print(config.SMPL_MODEL_DIR)
+        smplmodel = smplx.create(config.SMPL_MODEL_DIR,
+                                 model_type="smpl", gender="neutral", ext="pkl",
+                                 batch_size=self.batch_size).to(self.device)
+
+        # ## --- load the mean pose as original ----
+        smpl_mean_file = config.SMPL_MEAN_FILE
+
+        file = h5py.File(smpl_mean_file, 'r')
+        self.init_mean_pose = torch.from_numpy(file['pose'][:]).unsqueeze(0).repeat(self.batch_size, 1).float().to(self.device)
+        self.init_mean_shape = torch.from_numpy(file['shape'][:]).unsqueeze(0).repeat(self.batch_size, 1).float().to(self.device)
+        self.cam_trans_zero = torch.Tensor([0.0, 0.0, 0.0]).unsqueeze(0).to(self.device)
+        #
+
+        # # #-------------initialize SMPLify
+        self.smplify = SMPLify3D(smplxmodel=smplmodel,
+                            batch_size=self.batch_size,
+                            joints_category=self.joint_category,
+                            num_iters=self.num_smplify_iters,
+                            device=self.device)
+
+
+    def npy2smpl(self, npy_path):
+        out_path = npy_path.replace('.npy', '_rot.npy')
+        motions = np.load(npy_path, allow_pickle=True)[None][0]
+        # print_batch('', motions)
+        n_samples = motions['motion'].shape[0]
+        all_thetas = []
+        for sample_i in tqdm(range(n_samples)):
+            thetas, _ = self.joint2smpl(motions['motion'][sample_i].transpose(2, 0, 1))  # [nframes, njoints, 3]
+            all_thetas.append(thetas.cpu().numpy())
+        motions['motion'] = np.concatenate(all_thetas, axis=0)
+        print('motions', motions['motion'].shape)
+
+        print(f'Saving [{out_path}]')
+        np.save(out_path, motions)
+        exit()
+
+
+
+    def joint2smpl(self, input_joints, init_params=None):
+        _smplify = self.smplify # if init_params is None else self.smplify_fast
+        pred_pose = torch.zeros(self.batch_size, 72).to(self.device)
+        pred_betas = torch.zeros(self.batch_size, 10).to(self.device)
+        pred_cam_t = torch.zeros(self.batch_size, 3).to(self.device)
+        keypoints_3d = torch.zeros(self.batch_size, self.num_joints, 3).to(self.device)
+
+        # run the whole seqs
+        num_seqs = input_joints.shape[0]
+
+
+        # joints3d = input_joints[idx]  # *1.2 #scale problem [check first]
+        keypoints_3d = torch.Tensor(input_joints).to(self.device).float()
+
+        # if idx == 0:
+        if init_params is None:
+            pred_betas = self.init_mean_shape
+            pred_pose = self.init_mean_pose
+            pred_cam_t = self.cam_trans_zero
+        else:
+            pred_betas = init_params['betas']
+            pred_pose = init_params['pose']
+            pred_cam_t = init_params['cam']
+
+        if self.joint_category == "AMASS":
+            confidence_input = torch.ones(self.num_joints)
+            # make sure the foot and ankle
+            if self.fix_foot == True:
+                confidence_input[7] = 1.5
+                confidence_input[8] = 1.5
+                confidence_input[10] = 1.5
+                confidence_input[11] = 1.5
+        else:
+            print("Such category not settle down!")
+
+        new_opt_vertices, new_opt_joints, new_opt_pose, new_opt_betas, \
+        new_opt_cam_t, new_opt_joint_loss = _smplify(
+            pred_pose.detach(),
+            pred_betas.detach(),
+            pred_cam_t.detach(),
+            keypoints_3d,
+            conf_3d=confidence_input.to(self.device),
+            # seq_ind=idx
+        )
+
+        thetas = new_opt_pose.reshape(self.batch_size, 24, 3)
+        thetas = geometry.matrix_to_rotation_6d(geometry.axis_angle_to_matrix(thetas))  # [bs, 24, 6]
+        root_loc = torch.tensor(keypoints_3d[:, 0])  # [bs, 3]
+        root_loc = torch.cat([root_loc, torch.zeros_like(root_loc)], dim=-1).unsqueeze(1)  # [bs, 1, 6]
+        thetas = torch.cat([thetas, root_loc], dim=1).unsqueeze(0).permute(0, 2, 3, 1)  # [1, 25, 6, 196]
+
+        return thetas.clone().detach(), {'pose': new_opt_joints[0, :24].flatten().clone().detach(), 'betas': new_opt_betas.clone().detach(), 'cam': new_opt_cam_t.clone().detach()}
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_path", type=str, required=True, help='Blender file or dir with blender files')
+    parser.add_argument("--cuda", type=bool, default=True, help='')
+    parser.add_argument("--device", type=int, default=0, help='')
+    params = parser.parse_args()
+
+    simplify = joints2smpl(device_id=params.device, cuda=params.cuda)
+
+    if os.path.isfile(params.input_path) and params.input_path.endswith('.npy'):
+        simplify.npy2smpl(params.input_path)
+    elif os.path.isdir(params.input_path):
+        files = [os.path.join(params.input_path, f) for f in os.listdir(params.input_path) if f.endswith('.npy')]
+        for f in files:
+            simplify.npy2smpl(f)

visualize/vis_utils.py

@@ -0,0 +1,66 @@
+from model.rotation2xyz import Rotation2xyz
+import numpy as np
+from trimesh import Trimesh
+import os
+import torch
+from visualize.simplify_loc2rot import joints2smpl
+
+class npy2obj:
+    def __init__(self, npy_path, sample_idx, rep_idx, device=0, cuda=True):
+        self.npy_path = npy_path
+        self.motions = np.load(self.npy_path, allow_pickle=True)
+        if self.npy_path.endswith('.npz'):
+            self.motions = self.motions['arr_0']
+        self.motions = self.motions[None][0]
+        self.rot2xyz = Rotation2xyz(device='cpu')
+        self.faces = self.rot2xyz.smpl_model.faces
+        self.bs, self.njoints, self.nfeats, self.nframes = self.motions['motion'].shape
+        self.opt_cache = {}
+        self.sample_idx = sample_idx
+        self.total_num_samples = self.motions['num_samples']
+        self.rep_idx = rep_idx
+        self.absl_idx = self.rep_idx*self.total_num_samples + self.sample_idx
+        self.num_frames = self.motions['motion'][self.absl_idx].shape[-1]
+        self.j2s = joints2smpl(num_frames=self.num_frames, device_id=device, cuda=cuda)
+
+        if self.nfeats == 3:
+            print(f'Running SMPLify For sample [{sample_idx}], repetition [{rep_idx}], it may take a few minutes.')
+            motion_tensor, opt_dict = self.j2s.joint2smpl(self.motions['motion'][self.absl_idx].transpose(2, 0, 1))  # [nframes, njoints, 3]
+            self.motions['motion'] = motion_tensor.cpu().numpy()
+        elif self.nfeats == 6:
+            self.motions['motion'] = self.motions['motion'][[self.absl_idx]]
+        self.bs, self.njoints, self.nfeats, self.nframes = self.motions['motion'].shape
+        self.real_num_frames = self.motions['lengths'][self.absl_idx]
+
+        self.vertices = self.rot2xyz(torch.tensor(self.motions['motion']), mask=None,
+                                     pose_rep='rot6d', translation=True, glob=True,
+                                     jointstype='vertices',
+                                     # jointstype='smpl',  # for joint locations
+                                     vertstrans=True)
+        self.root_loc = self.motions['motion'][:, -1, :3, :].reshape(1, 1, 3, -1)
+        self.vertices += self.root_loc
+
+    def get_vertices(self, sample_i, frame_i):
+        return self.vertices[sample_i, :, :, frame_i].squeeze().tolist()
+
+    def get_trimesh(self, sample_i, frame_i):
+        return Trimesh(vertices=self.get_vertices(sample_i, frame_i),
+                       faces=self.faces)
+
+    def save_obj(self, save_path, frame_i):
+        mesh = self.get_trimesh(0, frame_i)
+        with open(save_path, 'w') as fw:
+            mesh.export(fw, 'obj')
+        return save_path
+    
+    def save_npy(self, save_path):
+        data_dict = {
+            'motion': self.motions['motion'][0, :, :, :self.real_num_frames],
+            'thetas': self.motions['motion'][0, :-1, :, :self.real_num_frames],
+            'root_translation': self.motions['motion'][0, -1, :3, :self.real_num_frames],
+            'faces': self.faces,
+            'vertices': self.vertices[0, :, :, :self.real_num_frames],
+            'text': self.motions['text'][0],
+            'length': self.real_num_frames,
+        }
+        np.save(save_path, data_dict)