Spaces:
Sleeping
Sleeping
| 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) | |