smpl_process.py

'''
@Desc: This file is used to process the smpl: convert rotation matrix to axis-angle
'''

import numpy as np
import os, os.path as osp, sys, pdb, argparse
from scipy.spatial.transform import Rotation

def convert_rotmat_to_axis_angle(rotmat):
    '''
    Args: rotmat: (N, 3, 3)
    Return: axis_angle: (N, 3)
    '''
    rotation = Rotation.from_matrix(rotmat)
    axis_angle = rotation.as_rotvec()
    return axis_angle

def parse_args():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--root', help='the root dir of an identity')
    argparser.add_argument('--smpl_file')
    return argparser.parse_args()

def main():
    args = parse_args()    
    root = args.root
    smpl_file = osp.join(root, args.smpl_file)
    npz = np.load(smpl_file, allow_pickle=True)
    smpl_params = {}
    for key in npz.keys():
        smpl_params[key] = npz[key]

    pdb.set_trace()
    rotmat_batch = smpl_params['poses'].reshape(-1, 3, 3)
    smpl_params['poses'] = convert_rotmat_to_axis_angle(rotmat_batch).reshape(-1, 24, 3)
    np.savez(smpl_file, **smpl_params)
    return

if __name__ == '__main__':
    main()