OpenMotionLab/MotionGPT

1 2import os3os.environ['DISPLAY'] = ':0.0'4os.environ['PYOPENGL_PLATFORM'] = 'osmesa'5os.environ["MUJOCO_GL"] = "osmesa"6from argparse import ArgumentParser7import numpy as np8import OpenGL.GL as gl9import imageio10import cv211import random12import torch13import moviepy.editor as mp14from scipy.spatial.transform import Rotation as RRR15import mGPT.render.matplot.plot_3d_global as plot_3d16from mGPT.render.pyrender.hybrik_loc2rot import HybrIKJointsToRotmat17from mGPT.render.pyrender.smpl_render import SMPLRender18 19 20if __name__ == '__main__':21 22    parser = ArgumentParser()23    parser.add_argument('--joints_path', type=str, help='Path to joints data')24    parser.add_argument('--method', type=str, help='Method for rendering')25    parser.add_argument('--output_mp4_path', type=str, help='Path to output MP4 file')26    parser.add_argument('--smpl_model_path', type=str, help='Path to SMPL model')27 28    args = parser.parse_args()29 30    joints_path = args.joints_path31    method = args.method32    output_mp4_path = args.output_mp4_path33    smpl_model_path = args.smpl_model_path34 35 36    data = np.load(joints_path)37 38    if method == 'slow':39        if len(data.shape) == 4:40            data = data[0]41        data = data - data[0, 0]42        pose_generator = HybrIKJointsToRotmat()43        pose = pose_generator(data)44        pose = np.concatenate([45            pose,46            np.stack([np.stack([np.eye(3)] * pose.shape[0], 0)] * 2, 1)47        ], 1)48        shape = [768, 768]49        render = SMPLRender(smpl_model_path)50 51        r = RRR.from_rotvec(np.array([np.pi, 0.0, 0.0]))52        pose[:, 0] = np.matmul(r.as_matrix().reshape(1, 3, 3), pose[:, 0])53        vid = []54        aroot = data[:, 0]55        aroot[:, 1:] = -aroot[:, 1:]56        params = dict(pred_shape=np.zeros([1, 10]),57                        pred_root=aroot,58                        pred_pose=pose)59        render.init_renderer([shape[0], shape[1], 3], params)60        for i in range(data.shape[0]):61            renderImg = render.render(i)62            vid.append(renderImg)63 64        out = np.stack(vid, axis=0)65        output_gif_path = output_mp4_path[:-4] + '.gif'66        imageio.mimwrite(output_gif_path, out, duration=50)67        out_video = mp.VideoFileClip(output_gif_path)68        out_video.write_videofile(output_mp4_path)69 70    elif method == 'fast':71        output_gif_path = output_mp4_path[:-4] + '.gif'72        if len(data.shape) == 3:73            data = data[None]74        if isinstance(data, torch.Tensor):75            data = data.cpu().numpy()76        pose_vis = plot_3d.draw_to_batch(data, [''], [output_gif_path])77        out_video = mp.VideoFileClip(output_gif_path)78        out_video.write_videofile(output_mp4_path)79