Spaces:
Running
Running
File size: 4,933 Bytes
517683d 3714210 517683d deda08a 517683d deda08a 517683d d1192f1 517683d d1192f1 517683d |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
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) |