Spaces:
Sleeping
Sleeping
File size: 2,747 Bytes
f0a0274 26344c4 3d5431e f0a0274 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
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)
del out, render
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)
del pose_vis
|