Spaces:
Running
on
L40S
Running
on
L40S
File size: 5,556 Bytes
38e20ed |
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 |
import os
import tempfile
import cv2
import kiui.mesh
import numpy as np
# os.environ['PYOPENGL_PLATFORM'] = 'osmesa' # osmesa or egl
os.environ['PYOPENGL_PLATFORM'] = 'egl'
import pyrender
import trimesh
# from psbody.mesh import Mesh
class MeshRenderer:
def __init__(self, size, fov=16 / 180 * np.pi, camera_pose=None, light_pose=None, black_bg=False):
# Camera
self.frustum = {'near': 0.01, 'far': 3.0}
self.camera = pyrender.PerspectiveCamera(yfov=fov, znear=self.frustum['near'],
zfar=self.frustum['far'], aspectRatio=1.0)
# Material
self.primitive_material = pyrender.material.MetallicRoughnessMaterial(
alphaMode='BLEND',
baseColorFactor=[0.3, 0.3, 0.3, 1.0],
metallicFactor=0.8,
roughnessFactor=0.8
)
# Lighting
light_color = np.array([1., 1., 1.])
self.light = pyrender.DirectionalLight(color=light_color, intensity=2)
self.light_angle = np.pi / 6.0
# Scene
self.scene = None
self._init_scene(black_bg)
# add camera and lighting
self._init_camera(camera_pose)
self._init_lighting(light_pose)
# Renderer
self.renderer = pyrender.OffscreenRenderer(*size, point_size=1.0)
def _init_scene(self, black_bg=False):
if black_bg:
self.scene = pyrender.Scene(ambient_light=[.2, .2, .2], bg_color=[0, 0, 0])
else:
self.scene = pyrender.Scene(ambient_light=[.2, .2, .2], bg_color=[255, 255, 255])
def _init_camera(self, camera_pose=None):
if camera_pose is None:
camera_pose = np.eye(4)
camera_pose[:3, 3] = np.array([0, 0, 1])
self.camera_pose = camera_pose.copy()
self.camera_node = self.scene.add(self.camera, pose=camera_pose)
def _init_lighting(self, light_pose=None):
if light_pose is None:
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, 0, 1])
self.light_pose = light_pose.copy()
light_poses = self._get_light_poses(self.light_angle, light_pose)
self.light_nodes = [self.scene.add(self.light, pose=light_pose) for light_pose in light_poses]
def set_camera_pose(self, camera_pose):
self.camera_pose = camera_pose.copy()
self.scene.set_pose(self.camera_node, pose=camera_pose)
def set_lighting_pose(self, light_pose):
self.light_pose = light_pose.copy()
light_poses = self._get_light_poses(self.light_angle, light_pose)
for light_node, light_pose in zip(self.light_nodes, light_poses):
self.scene.set_pose(light_node, pose=light_pose)
def render_mesh(self, v, f, t_center, rot=np.zeros(3), tex_img=None, tex_uv=None,
camera_pose=None, light_pose=None):
# Prepare mesh
v[:] = cv2.Rodrigues(rot)[0].dot((v - t_center).T).T + t_center
if tex_img is not None:
tex = pyrender.Texture(source=tex_img, source_channels='RGB')
tex_material = pyrender.material.MetallicRoughnessMaterial(baseColorTexture=tex)
from kiui.mesh import Mesh
import torch
mesh = Mesh(
v=torch.from_numpy(v),
f=torch.from_numpy(f),
vt=tex_uv['vt'],
ft=tex_uv['ft']
)
with tempfile.NamedTemporaryFile(suffix='.obj') as f:
mesh.write_obj(f.name)
tri_mesh = trimesh.load(f.name, process=False)
return tri_mesh
# tri_mesh = self._pyrender_mesh_workaround(mesh)
render_mesh = pyrender.Mesh.from_trimesh(tri_mesh, material=tex_material)
else:
tri_mesh = trimesh.Trimesh(vertices=v, faces=f)
render_mesh = pyrender.Mesh.from_trimesh(tri_mesh, material=self.primitive_material, smooth=True)
mesh_node = self.scene.add(render_mesh, pose=np.eye(4))
# Change camera and lighting pose if necessary
if camera_pose is not None:
self.set_camera_pose(camera_pose)
if light_pose is not None:
self.set_lighting_pose(light_pose)
# Render
flags = pyrender.RenderFlags.SKIP_CULL_FACES
color, depth = self.renderer.render(self.scene, flags=flags)
# Remove mesh
self.scene.remove_node(mesh_node)
return color, depth
@staticmethod
def _get_light_poses(light_angle, light_pose):
light_poses = []
init_pos = light_pose[:3, 3].copy()
light_poses.append(light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([light_angle, 0, 0]))[0].dot(init_pos)
light_poses.append(light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([-light_angle, 0, 0]))[0].dot(init_pos)
light_poses.append(light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([0, -light_angle, 0]))[0].dot(init_pos)
light_poses.append(light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([0, light_angle, 0]))[0].dot(init_pos)
light_poses.append(light_pose.copy())
return light_poses
@staticmethod
def _pyrender_mesh_workaround(mesh):
# Workaround as pyrender requires number of vertices and uv coordinates to be the same
with tempfile.NamedTemporaryFile(suffix='.obj') as f:
mesh.write_obj(f.name)
tri_mesh = trimesh.load(f.name, process=False)
return tri_mesh
|