Spaces:
Running
Running
# adapted from HuMoR | |
import torch | |
import numpy as np | |
from tqdm import tqdm | |
import trimesh | |
from .parameters import colors, smpl_connections | |
from .mesh_viewer import MeshViewer | |
c2c = lambda tensor: tensor.detach().cpu().numpy() # noqa | |
def viz_smpl_seq( | |
pyrender, | |
out_path, | |
body, | |
# | |
start=None, | |
end=None, | |
# | |
imw=720, | |
imh=720, | |
fps=20, | |
use_offscreen=True, | |
follow_camera=True, | |
progress_bar=tqdm, | |
# | |
contacts=None, | |
render_body=True, | |
render_joints=False, | |
render_skeleton=False, | |
render_ground=True, | |
ground_plane=None, | |
wireframe=False, | |
RGBA=False, | |
joints_seq=None, | |
joints_vel=None, | |
vtx_list=None, | |
points_seq=None, | |
points_vel=None, | |
static_meshes=None, | |
camera_intrinsics=None, | |
img_seq=None, | |
point_rad=0.015, | |
skel_connections=smpl_connections, | |
img_extn="png", | |
ground_alpha=1.0, | |
body_alpha=None, | |
mask_seq=None, | |
cam_offset=[0.0, 2.2, 0.9], # [0.0, 4.0, 1.25], | |
ground_color0=[0.8, 0.9, 0.9], | |
ground_color1=[0.6, 0.7, 0.7], | |
# ground_color0=[1.0, 1.0, 1.0], | |
# ground_color1=[0.0, 0.0, 0.0], | |
skel_color=[0.5, 0.5, 0.5], # [0.0, 0.0, 1.0], | |
joint_rad=0.015, | |
point_color=[0.0, 0.0, 1.0], | |
joint_color=[0.0, 1.0, 0.0], | |
contact_color=[1.0, 0.0, 0.0], | |
vertex_color=[], | |
vertex_color_list = [], | |
# color_vtx0 = colors['neon pink'] | |
# color_vtx1 = colors['sky blue'] | |
render_bodies_static=None, | |
render_points_static=None, | |
cam_rot=None, | |
): | |
""" | |
Visualizes the body model output of a smpl sequence. | |
- body : body model output from SMPL forward pass (where the sequence is the batch) | |
- joints_seq : list of torch/numy tensors/arrays | |
- points_seq : list of torch/numpy tensors | |
- camera_intrinsics : (fx, fy, cx, cy) | |
- ground_plane : [a, b, c, d] | |
- render_bodies_static is an integer, if given renders all bodies at once but only every x steps | |
""" | |
if contacts is not None and torch.is_tensor(contacts): | |
contacts = c2c(contacts) | |
# import ipdb; ipdb.set_trace() | |
if isinstance(body, list): | |
render_pair = True | |
else: | |
render_pair = False | |
if render_body or vtx_list is not None: | |
if render_pair: | |
nv = body[0].v.size(1) | |
else: | |
nv = body.v.size(1) | |
vertex_colors = np.tile(vertex_color, (nv, 1)) | |
if body_alpha is not None: | |
vtx_alpha = np.ones((vertex_colors.shape[0], 1)) * body_alpha | |
vertex_colors = np.concatenate([vertex_colors, vtx_alpha], axis=1) | |
if render_pair: | |
faces = c2c(body[0].f) | |
else: | |
faces = c2c(body.f) | |
if not vertex_color_list: | |
vertex_color_list= ['neon pink','sky blue'] | |
if not vertex_color: | |
vertex_color=colors['sky blue'] | |
else: | |
vertex_color=colors[vertex_color[0]] | |
if render_pair: | |
color_vtx0=colors[vertex_color_list[0]] | |
color_vtx1=colors[vertex_color_list[1]] | |
body_mesh_seq0 = [ | |
trimesh.Trimesh( | |
vertices=c2c(body[0].v[i]), | |
faces=faces, | |
vertex_colors=color_vtx0, | |
process=False, | |
) | |
for i in range(body[0].v.size(0)) | |
] | |
body_mesh_seq1 = [ | |
trimesh.Trimesh( | |
vertices=c2c(body[1].v[i]), | |
faces=faces, | |
vertex_colors=color_vtx1, | |
process=False, | |
) | |
for i in range(body[1].v.size(0)) | |
] | |
else: | |
body_mesh_seq = [ | |
trimesh.Trimesh( | |
vertices=c2c(body.v[i]), | |
faces=faces, | |
vertex_colors=vertex_color, | |
process=False, | |
) | |
for i in range(body.v.size(0)) | |
] | |
if render_joints and joints_seq is None: | |
# only body joints | |
joints_seq = [c2c(body.Jtr[i, :22]) for i in range(body.Jtr.size(0))] | |
elif render_joints and torch.is_tensor(joints_seq[0]): | |
joints_seq = [c2c(joint_frame) for joint_frame in joints_seq] | |
if joints_vel is not None and torch.is_tensor(joints_vel[0]): | |
joints_vel = [c2c(joint_frame) for joint_frame in joints_vel] | |
if points_vel is not None and torch.is_tensor(points_vel[0]): | |
points_vel = [c2c(joint_frame) for joint_frame in points_vel] | |
# cam_offset = [0.0, 2.2, 0.2] | |
mv = MeshViewer( | |
pyrender, | |
width=imw, | |
height=imh, | |
use_offscreen=use_offscreen, | |
follow_camera=follow_camera, | |
camera_intrinsics=camera_intrinsics, | |
img_extn=img_extn, | |
default_cam_offset=cam_offset, | |
default_cam_rot=cam_rot, | |
) | |
if render_body and render_bodies_static is None: | |
if render_pair: | |
mv.add_mesh_seq(body_mesh_seq0, progress_bar=progress_bar) | |
mv.add_mesh_seq(body_mesh_seq1, progress_bar=progress_bar) | |
else: | |
mv.add_mesh_seq(body_mesh_seq, progress_bar=progress_bar) | |
elif render_body and render_bodies_static is not None: | |
if render_pair: | |
mv.add_static_meshes( | |
[ | |
body_mesh_seq0[i] | |
for i in range(len(body_mesh_seq0)) | |
if i % render_bodies_static == 0 | |
] | |
) | |
mv.add_static_meshes( | |
[ | |
body_mesh_seq1[i] | |
for i in range(len(body_mesh_seq1)) | |
if i % render_bodies_static == 0 | |
] | |
) | |
else: | |
mv.add_static_meshes( | |
[ | |
body_mesh_seq[i] | |
for i in range(len(body_mesh_seq)) | |
if i % render_bodies_static == 0 | |
] | |
) | |
if render_joints and render_skeleton: | |
mv.add_point_seq( | |
joints_seq, | |
color=joint_color, | |
radius=joint_rad, | |
contact_seq=contacts, | |
connections=skel_connections, | |
connect_color=skel_color, | |
vel=joints_vel, | |
contact_color=contact_color, | |
render_static=render_points_static, | |
) | |
elif render_joints: | |
mv.add_point_seq( | |
joints_seq, | |
color=joint_color, | |
radius=joint_rad, | |
contact_seq=contacts, | |
vel=joints_vel, | |
contact_color=contact_color, | |
render_static=render_points_static, | |
) | |
if vtx_list is not None: | |
mv.add_smpl_vtx_list_seq( | |
body_mesh_seq, vtx_list, color=[0.0, 0.0, 1.0], radius=0.015 | |
) | |
if points_seq is not None: | |
if torch.is_tensor(points_seq[0]): | |
points_seq = [c2c(point_frame) for point_frame in points_seq] | |
mv.add_point_seq( | |
points_seq, | |
color=point_color, | |
radius=point_rad, | |
vel=points_vel, | |
render_static=render_points_static, | |
) | |
if static_meshes is not None: | |
mv.set_static_meshes(static_meshes) | |
if img_seq is not None: | |
mv.set_img_seq(img_seq) | |
if mask_seq is not None: | |
mv.set_mask_seq(mask_seq) | |
if render_ground: | |
xyz_orig = None | |
if ground_plane is not None: | |
if render_body: | |
if render_pair: | |
xyz_orig = (body_mesh_seq0[0].vertices[0, :] + body_mesh_seq1[0].vertices[0, :]) / 2 | |
else: | |
xyz_orig = body_mesh_seq[0].vertices[0, :] | |
elif render_joints: | |
xyz_orig = joints_seq[0][0, :] | |
elif points_seq is not None: | |
xyz_orig = points_seq[0][0, :] | |
mv.add_ground( | |
ground_plane=ground_plane, | |
xyz_orig=xyz_orig, | |
color0=ground_color0, | |
color1=ground_color1, | |
alpha=ground_alpha, | |
) | |
mv.set_render_settings( | |
out_path=out_path, | |
wireframe=wireframe, | |
RGBA=RGBA, | |
single_frame=( | |
render_points_static is not None or render_bodies_static is not None | |
), | |
) # only does anything for offscreen rendering | |
try: | |
mv.animate(fps=fps, start=start, end=end, progress_bar=progress_bar) | |
except RuntimeError as err: | |
print("Could not render properly with the error: %s" % (str(err))) | |
del mv |