|
|
| import os |
| os.environ['DISPLAY'] = ':0.0' |
| os.environ['PYOPENGL_PLATFORM'] = 'osmesa' |
| os.environ["MUJOCO_GL"] = "osmesa" |
| from argparse import ArgumentParser |
| import numpy as np |
| import OpenGL.GL as gl |
| import imageio |
| import cv2 |
| import random |
| import torch |
| import moviepy.editor as mp |
| from scipy.spatial.transform import Rotation as RRR |
| import mGPT.render.matplot.plot_3d_global as plot_3d |
| from mGPT.render.pyrender.hybrik_loc2rot import HybrIKJointsToRotmat |
| from mGPT.render.pyrender.smpl_render import SMPLRender |
|
|
|
|
| if __name__ == '__main__': |
|
|
| parser = ArgumentParser() |
| parser.add_argument('--joints_path', type=str, help='Path to joints data') |
| parser.add_argument('--method', type=str, help='Method for rendering') |
| parser.add_argument('--output_mp4_path', type=str, help='Path to output MP4 file') |
| parser.add_argument('--smpl_model_path', type=str, help='Path to SMPL model') |
|
|
| args = parser.parse_args() |
|
|
| joints_path = args.joints_path |
| method = args.method |
| output_mp4_path = args.output_mp4_path |
| smpl_model_path = args.smpl_model_path |
|
|
|
|
| data = np.load(joints_path) |
|
|
| if method == 'slow': |
| if len(data.shape) == 4: |
| data = data[0] |
| data = data - data[0, 0] |
| pose_generator = HybrIKJointsToRotmat() |
| pose = pose_generator(data) |
| pose = np.concatenate([ |
| pose, |
| np.stack([np.stack([np.eye(3)] * pose.shape[0], 0)] * 2, 1) |
| ], 1) |
| shape = [768, 768] |
| render = SMPLRender(smpl_model_path) |
|
|
| r = RRR.from_rotvec(np.array([np.pi, 0.0, 0.0])) |
| pose[:, 0] = np.matmul(r.as_matrix().reshape(1, 3, 3), pose[:, 0]) |
| vid = [] |
| aroot = data[:, 0] |
| aroot[:, 1:] = -aroot[:, 1:] |
| params = dict(pred_shape=np.zeros([1, 10]), |
| pred_root=aroot, |
| pred_pose=pose) |
| render.init_renderer([shape[0], shape[1], 3], params) |
| for i in range(data.shape[0]): |
| renderImg = render.render(i) |
| vid.append(renderImg) |
|
|
| out = np.stack(vid, axis=0) |
| output_gif_path = output_mp4_path[:-4] + '.gif' |
| imageio.mimwrite(output_gif_path, out, duration=50) |
| out_video = mp.VideoFileClip(output_gif_path) |
| out_video.write_videofile(output_mp4_path) |
|
|
| elif method == 'fast': |
| output_gif_path = output_mp4_path[:-4] + '.gif' |
| if len(data.shape) == 3: |
| data = data[None] |
| if isinstance(data, torch.Tensor): |
| data = data.cpu().numpy() |
| pose_vis = plot_3d.draw_to_batch(data, [''], [output_gif_path]) |
| out_video = mp.VideoFileClip(output_gif_path) |
| out_video.write_videofile(output_mp4_path) |
|
|
