Spaces:
Running
Running
import numpy as np | |
import torch | |
from .humor_render_tools.tools import viz_smpl_seq | |
from smplx.utils import Struct | |
from .video import Video | |
import os | |
from multiprocessing import Pool | |
from tqdm import tqdm | |
from multiprocessing import Process | |
THIS_FOLDER = os.path.dirname(os.path.abspath(__file__)) | |
FACE_PATH = os.path.join(THIS_FOLDER, "humor_render_tools/smplh.faces") | |
FACES = torch.from_numpy(np.int32(np.load(FACE_PATH))) | |
class HumorRenderer: | |
def __init__(self, fps=30.0, **kwargs): | |
self.kwargs = kwargs | |
self.fps = fps | |
def __call__(self, vertices, output, render_pair=False,text=None,colors=[],**kwargs): | |
params = self.kwargs | kwargs | |
fps = self.fps | |
if "fps" in params: | |
fps = params.pop("fps") | |
if render_pair: | |
render_overlaid(vertices, output,fps,colors,**params) | |
else: | |
render(vertices, output, fps,colors,**params) | |
fname = f'{output}.mp4' | |
return fname | |
def render_overlaid(vertices, out_path, fps,colors=[], progress_bar=tqdm,**kwargs): | |
assert isinstance(vertices, list) and len(vertices) == 2 | |
# Put the vertices at the floor level | |
# F X N x 3 ===> [F X N x 3, F X N x 3] | |
ground = vertices[0][..., 2].min() | |
vertices[0][..., 2] -= ground | |
vertices[1][..., 2] -= ground | |
verts0 = vertices[0] | |
verts1 = vertices[1] | |
import pyrender | |
# remove title if it exists | |
kwargs.pop("title", None) | |
# vertices: SMPL-H vertices | |
# verts = np.load("interval_2_verts.npy") | |
out_folder = os.path.splitext(out_path)[0] | |
verts0 = torch.from_numpy(verts0) | |
body_pred0 = Struct(v=verts0, f=FACES) | |
verts1 = torch.from_numpy(verts1) | |
body_pred1 = Struct(v=verts1, f=FACES) | |
# out_folder, body_pred, start, end, fps, kwargs = args | |
viz_smpl_seq( | |
pyrender, out_folder, [body_pred0, body_pred1], fps=fps,progress_bar=progress_bar,vertex_color_list=colors, **kwargs | |
) | |
video = Video(out_folder, fps=fps) | |
video.save(out_path) | |
def render(vertices, out_path, fps, colors=[], progress_bar=tqdm, **kwargs): | |
# Put the vertices at the floor level | |
ground = vertices[..., 2].min() | |
vertices[..., 2] -= ground | |
import pyrender | |
# remove title if it exists | |
kwargs.pop("title", None) | |
# vertices: SMPL-H vertices | |
# verts = np.load("interval_2_verts.npy") | |
out_folder = os.path.splitext(out_path)[0] | |
verts = torch.from_numpy(vertices) | |
body_pred = Struct(v=verts, f=FACES) | |
# out_folder, body_pred, start, end, fps, kwargs = args | |
viz_smpl_seq( | |
pyrender, out_folder, body_pred, fps=fps,progress_bar=progress_bar,vertex_color=colors, **kwargs | |
) | |
video = Video(out_folder, fps=fps) | |
video.save(out_path) | |
import shutil | |
shutil.rmtree(out_folder) | |
def render_offset(args): | |
import pyrender | |
out_folder, body_pred, start, end, fps, kwargs = args | |
viz_smpl_seq( | |
pyrender, out_folder, body_pred, start=start, end=end, fps=fps, **kwargs | |
) | |
return 0 | |
def render_multiprocess(vertices, out_path, fps, **kwargs): | |
# WIP: does not work yet | |
# import ipdb | |
# ipdb.set_trace() | |
# remove title if it exists | |
kwargs.pop("title", None) | |
# vertices: SMPL-H vertices | |
# verts = np.load("interval_2_verts.npy") | |
out_folder = os.path.splitext(out_path)[0] | |
verts = torch.from_numpy(vertices) | |
body_pred = Struct(v=verts, f=FACES) | |
# faster rendering | |
# by rendering part of the sequence in parallel | |
# still work in progress, use one process for now | |
n_processes = 1 | |
verts_lst = np.array_split(verts, n_processes) | |
len_split = [len(x) for x in verts_lst] | |
starts = [0] + np.cumsum([x for x in len_split[:-1]]).tolist() | |
ends = np.cumsum([x for x in len_split]).tolist() | |
out_folders = [out_folder for _ in range(n_processes)] | |
fps_s = [fps for _ in range(n_processes)] | |
kwargs_s = [kwargs for _ in range(n_processes)] | |
body_pred_s = [body_pred for _ in range(n_processes)] | |
arguments = [out_folders, body_pred_s, starts, ends, fps_s, kwargs_s] | |
# sanity | |
# lst = [verts[start:end] for start, end in zip(starts, ends)] | |
# assert (torch.cat(lst) == verts).all() | |
processes = [] | |
for _, args in zip(range(n_processes), zip(*arguments)): | |
process = Process(target=render_offset, args=(args,)) | |
process.start() | |
processes.append(process) | |
for process in processes: | |
process.join() | |
if False: | |
# start 4 worker processes | |
with Pool(processes=n_processes) as pool: | |
# print "[0, 1, 4,..., 81]" | |
# print same numbers in arbitrary order | |
print(f"0/{n_processes} rendered") | |
i = 0 | |
for _ in pool.imap_unordered(render_offset, zip(*arguments)): | |
i += 1 | |
print(f"i/{n_processes} rendered") | |
video = Video(out_folder, fps=fps) | |
video.save(out_path) |