atnikos's picture
fix rendering and args in textboxes
deda08a
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)