File size: 5,724 Bytes
4409449 |
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 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
"""
This script is borrowed from https://github.com/mkocabas/VIBE
Adhere to their licence to use this script
It has been modified
"""
import os
import math
import trimesh
import pyrender
import numpy as np
from pyrender.constants import RenderFlags
# os.environ['DISPLAY'] = ':0.0'
# os.environ['PYOPENGL_PLATFORM'] = 'egl'
# os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
SMPL_MODEL_DIR = "data/smpl_data/"
def get_smpl_faces():
return np.load(os.path.join(SMPL_MODEL_DIR, "smplfaces.npy"))
class WeakPerspectiveCamera(pyrender.Camera):
def __init__(self,
scale,
translation,
znear=pyrender.camera.DEFAULT_Z_NEAR,
zfar=None,
name=None):
super(WeakPerspectiveCamera, self).__init__(
znear=znear,
zfar=zfar,
name=name,
)
self.scale = scale
self.translation = translation
def get_projection_matrix(self, width=None, height=None):
P = np.eye(4)
P[0, 0] = self.scale[0]
P[1, 1] = self.scale[1]
P[0, 3] = self.translation[0] * self.scale[0]
P[1, 3] = -self.translation[1] * self.scale[1]
P[2, 2] = -1
return P
class Renderer:
def __init__(self, background=None, resolution=(224, 224), bg_color=[0, 0, 0, 0.5], orig_img=False, wireframe=False, cam_pose=np.eye(4)):
width, height = resolution
self.background = np.zeros((height, width, 3))
self.resolution = resolution
self.faces = get_smpl_faces()
self.orig_img = orig_img
self.wireframe = wireframe
self.renderer = pyrender.OffscreenRenderer(
viewport_width=self.resolution[0],
viewport_height=self.resolution[1],
point_size=0.5
)
# set the scene
self.scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4))
light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=4)
light_pose = np.eye(4)
light_pose[:3, 3] = [0, -1, 1]
self.scene.add(light, pose=np.dot(cam_pose,light_pose).copy())
light_pose[:3, 3] = [0, 1, 1]
self.scene.add(light, pose=np.dot(cam_pose,light_pose).copy())
light_pose[:3, 3] = [1, 1, 2]
self.scene.add(light, pose=np.dot(cam_pose,light_pose).copy())
"""ok
light_pose = np.eye(4)
light_pose[:3, 3] = [0, -1, 1]
self.scene.add(light, pose=light_pose)
light_pose[:3, 3] = [0, 1, 1]
self.scene.add(light, pose=light_pose)
light_pose[:3, 3] = [1, 1, 2]
self.scene.add(light, pose=light_pose)
"""
# light_pose[:3, 3] = [0, -2, 2]
# [droite, hauteur, profondeur camera]
"""
light_pose = np.eye(4)
light_pose[:3, 3] = [0, -1, 1]
self.scene.add(light, pose=light_pose)
light_pose[:3, 3] = [0, 1, 1]
self.scene.add(light, pose=light_pose)
light_pose[:3, 3] = [1, 1, 2]
self.scene.add(light, pose=light_pose)
"""
def render(self, img, verts, cam, angle=None, axis=None, mesh_filename=None, color=[1.0, 1.0, 0.9],
cam_pose=np.eye(4)):
mesh = trimesh.Trimesh(vertices=verts, faces=self.faces, process=False)
Rx = trimesh.transformations.rotation_matrix(math.radians(180), [1, 0, 0])
# Rx = trimesh.transformations.rotation_matrix(math.radians(-90), [1, 0, 0])
mesh.apply_transform(Rx)
if mesh_filename is not None:
mesh.export(mesh_filename)
if angle and axis:
R = trimesh.transformations.rotation_matrix(math.radians(angle), axis)
mesh.apply_transform(R)
sx, sy, tx, ty = cam
camera = WeakPerspectiveCamera(
scale=[sx, sy],
translation=[tx, ty],
zfar=100000.
)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0, # 0.0 for no specular lighting
# metallicFactor=0.7, # 0.0 for no specular lighting
alphaMode='OPAQUE',
baseColorFactor=(color[0], color[1], color[2], 1.0)
)
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
mesh_node = self.scene.add(mesh, 'mesh')
cam_node = self.scene.add(camera, pose=cam_pose)
if self.wireframe:
render_flags = RenderFlags.RGBA | RenderFlags.ALL_WIREFRAME
else:
render_flags = RenderFlags.RGBA
rgb, _ = self.renderer.render(self.scene, flags=render_flags)
if rgb.shape[-1]==3:
# Debug
# 0 not distinguish alpha
valid_mask = (rgb[:, :, -1] > 0)[:, :, np.newaxis]
output_img = rgb * valid_mask + (1 - valid_mask) * img
elif rgb.shape[-1]==4:
# valid_mask = (rgb[:, :, -1] > 128)[:, :, np.newaxis]
# output_img = rgb[:, :, :-1] * valid_mask + (1 - valid_mask) * img
# # output alpha
valid_mask = (rgb[:, :, -1] > 128)[:, :]
output_img = np.copy(rgb)
output_img[:, :, -1] *= valid_mask
# output_img = img
else:
raise ValueError(f"rgb shape {rgb.shape[-1]} is not correct!")
image = output_img.astype(np.uint8)
self.scene.remove_node(mesh_node)
self.scene.remove_node(cam_node)
return image
def get_renderer(width, height, cam_pose):
renderer = Renderer(resolution=(width, height),
bg_color=[1, 1, 1, 0.5],
orig_img=False,
wireframe=False,
cam_pose=cam_pose)
return renderer
|