Render animations in AMASS format to USD in Maya

# import packages
import numpy as np
from tqdm.auto import tqdm

# import GenMotion modules

from genmotion.dataset.amass_params import SMPL_H_SKELETON # recognize the skeleton type as SMPL_H
from genmotion.render.maya.utils import MayaController

# Set up Maya Controller (socket client) with PORT number
mc = MayaController(PORT=12345)

# Test Maya command by setting the key frame at timeline
mc.SetCurrentTimeFrame(0)

mc.SendPythonCommand("import maya.cmds as cmds")
mc.SendPythonCommand("import maya.api.OpenMaya as om")

'\n\x00'

def get_root_transform_matrix(root_name = "f_avg_root"):
    # https://help.autodesk.com/view/MAYAUL/2017/CHS/?guid=__py_ref_class_open_maya_1_1_m_transformation_matrix_html
    multi_line_cmd = f"""
# An object of interest:
cmds.currentTime(0, edit=True )
object = '{root_name}'
# Get the transform matrix as a list of 16 floats
m_list = cmds.xform(object, query=True, matrix=True)
# Create the MMatrix object
m = om.MMatrix(m_list)
m_inverse = m.inverse()
""".replace("\n",r"\n")

    # check success
    mc.SendPythonCommand("print(m, m_inverse)")
    mc.SendPythonCommand(multi_line_cmd)

def rotate_root_transform(joint_name="f_avg_root"):
    # https://help.autodesk.com/view/MAYAUL/2017/CHS/?guid=__py_ref_class_open_maya_1_1_m_transformation_matrix_html
    multi_line_cmd = f"""
# An object of interest:
joint = '{joint_name}'
# Get the transform matrix as a list of 16 floats
j_list = cmds.xform(joint, query=True, matrix=True)
# Create the MMatrix object
joint_m = om.MMatrix(j_list)
new_joint_m = joint_m * m_inverse
# rotate matrix
mt = om.MTransformationMatrix(new_joint_m)

# Get the rotations
rot = mt.rotation()
x,y,z = om.MAngle(rot.x).asDegrees(), om.MAngle(rot.y).asDegrees(), om.MAngle(rot.z).asDegrees()
cmds.setAttr('{joint_name}.rotateX', x)
cmds.setAttr('{joint_name}.rotateY', y)
cmds.setAttr('{joint_name}.rotateZ', z)

# Get the rotations
tran = mt.translation(om.MSpace.kWorld)
x,y,z = tran.x, tran.y, tran.z
cmds.setAttr('{joint_name}.translateX', x)
cmds.setAttr('{joint_name}.translateY', y)
cmds.setAttr('{joint_name}.translateZ', z)

""".replace("\n",r"\n")

    # check success
    mc.SendPythonCommand(multi_line_cmd)

# rotate_root_transform()

# get_root_transform_matrix()

# after downloading the dataset, unzip the tar.gz file, and assign the path for `npy` relative to the\
# the root folder of `GenMotion/src`
amass_npz_fname = 'D:/research/data/DFaust/DFaust_67/50002/50002_chicken_wings_poses.npz'

# load body data from npz file
bdata = np.load(amass_npz_fname)

# Get the frame rate of from the body data
print("framerate:", bdata['mocap_framerate'])

framerate: 60

# Set one frame for the animation
def SetOneFrame(frame, body_data, joints_info_list: list, joint_prefix = "f_avg_", axis_up=True):
    '''Set one maya frame from load data

    :param frame: int
    :param body_data: array
    :param joints_info_list: list
    :return:
    '''
    #  set current time frame
    mc.SetCurrentTimeFrame(frame)

    for i in range(0,len(joints_info_list)):
        # get rotation from body data
        rotateX, rotateY, rotateZ = np.rad2deg(body_data['poses'][frame][i*3:(i+1)*3])

        # get joint name
        joint_name = joint_prefix + joints_info_list[i]

        # if i == 0, set root translation
        if i == 0:
            joint_name = joint_prefix + "root"
            translateX,translateY,translateZ = body_data['trans'][frame]
            mc.SetObjectAttribute(joint_name, "translateX", translateX)
            mc.SetObjectAttribute(joint_name, "translateY", translateY)
            mc.SetObjectAttribute(joint_name, "translateZ", translateZ)

        # set rotation for joints
        mc.SetObjectAttribute(joint_name, "rotateX", rotateX)
        mc.SetObjectAttribute(joint_name, "rotateY", rotateY)
        mc.SetObjectAttribute(joint_name, "rotateZ", rotateZ)

        if axis_up and i == 0:
            joint_name = joint_prefix + "root"

            # if start, record initial root transform matrix
            if frame == 0:
                get_root_transform_matrix()

            rotate_root_transform()

        mc.SetCurrentKeyFrameForPositionAndRotation(joint_name)
        #break

# Now we maker animation for every 10 frame
frame_interval = 10
for frame in tqdm(range(0, len(bdata["poses"]), frame_interval)):
    SetOneFrame(frame, bdata, SMPL_H_SKELETON, axis_up=True)

# mc.SetNewScene()