''' @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()