body-pose-animation/utils/render.py
2021-02-08 15:33:27 +01:00

98 lines
2.9 KiB
Python

from typing import List, Set, Dict, Tuple, Optional
import numpy as np
import trimesh
import pyrender
def render_model(
scene,
model,
model_out,
color=[0.3, 0.3, 0.3, 0.8],
name=None,
replace=False,
pose=None
):
vertices = model_out.vertices.detach().cpu().numpy().squeeze()
# set vertex colors, maybe use this to highlight accuracies
vertex_colors = np.ones([vertices.shape[0], 4]) * color
# triangulate vertex mesh
tri_mesh = trimesh.Trimesh(vertices, model.faces,
vertex_colors=vertex_colors)
mesh = pyrender.Mesh.from_trimesh(tri_mesh)
if name is not None and replace:
for node in scene.get_nodes(name=name):
scene.remove_node(node)
return scene.add(mesh, name=name, pose=pose)
def render_model_with_tfs(
scene,
model,
model_out,
color=[0.3, 0.3, 0.3, 0.8],
name=None,
replace=False,
pose=None,
transforms=None
):
vertices = model_out.vertices.detach().cpu().numpy().squeeze()
# set vertex colors, maybe use this to highlight accuracies
vertex_colors = np.ones([vertices.shape[0], 4]) * color
# triangulate vertex mesh
tri_mesh = trimesh.Trimesh(vertices, model.faces,
vertex_colors=vertex_colors)
mesh = pyrender.Mesh.from_trimesh(tri_mesh, poses=transforms)
if name is not None and replace:
for node in scene.get_nodes(name=name):
scene.remove_node(node)
return scene.add(mesh, name=name, pose=pose)
def render_points(scene, points, radius=0.005, color=[0.0, 0.0, 1.0, 1.0], name=None, transform=None):
sm = trimesh.creation.uv_sphere(radius=radius)
sm.visual.vertex_colors = color
tfs = np.tile(np.eye(4), (len(points), 1, 1))
tfs[:, :3, 3] = points
if transform is not None:
tfs = transform @ tfs
pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
# return the render scene node
return scene.add(pcl, name=name)
def render_camera(scene, radius=0.5, height=0.5, color=[0.0, 0.0, 1.0, 1.0], name=None):
sm = trimesh.creation.cone(radius, height, sections=None, transform=None)
sm.visual.vertex_colors = color
tfs = np.eye(4)
pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
# return the render scene node
return scene.add(pcl, name=name)
def render_image_plane(scene, image, scale, name=None):
height, width, _ = image.shape
mat = trimesh.visual.texture.TextureVisuals(
image=image, uv=[[0, 0], [0, 1], [1, 0], [1, 1]])
tm = trimesh.load('plane.obj', visual=mat)
tm.visual = mat
tfs = np.eye(4)
tfs[0, 0] = width / height * scale
tfs[1, 1] *= scale
tfs[2, 2] *= scale
tfs[0, 3] = (width / height - 1) * scale
material2 = pyrender.Material(name=name, emissiveTexture=image)
m = pyrender.Mesh.from_trimesh(tm, poses=tfs)
return scene.add(m, name=name)