294 lines
12 KiB
Python
294 lines
12 KiB
Python
import os
|
||
import numpy as np
|
||
import cv2
|
||
import pyrender
|
||
import trimesh
|
||
import copy
|
||
# 这个顺序是BGR的。虽然render的使用的是RGB的，但是由于和图像拼接了，所以又变成BGR的了
|
||
colors = [
|
||
# (0.5, 0.2, 0.2, 1.), # Defalut BGR
|
||
(.5, .5, .7, 1.), # Pink BGR
|
||
(.44, .50, .98, 1.), # Red
|
||
(.7, .7, .6, 1.), # Neutral
|
||
(.5, .5, .7, 1.), # Blue
|
||
(.5, .55, .3, 1.), # capsule
|
||
(.3, .5, .55, 1.), # Yellow
|
||
# (.6, .6, .6, 1.), # gray
|
||
(.9, 1., 1., 1.),
|
||
(0.95, 0.74, 0.65, 1.),
|
||
(.9, .7, .7, 1.)
|
||
]
|
||
|
||
colors_table = {
|
||
# colorblind/print/copy safe:
|
||
'_blue': [0.65098039, 0.74117647, 0.85882353],
|
||
'_pink': [.9, .7, .7],
|
||
'_mint': [ 166/255., 229/255., 204/255.],
|
||
'_mint2': [ 202/255., 229/255., 223/255.],
|
||
'_green': [ 153/255., 216/255., 201/255.],
|
||
'_green2': [ 171/255., 221/255., 164/255.],
|
||
'_red': [ 251/255., 128/255., 114/255.],
|
||
'_orange': [ 253/255., 174/255., 97/255.],
|
||
'_yellow': [ 250/255., 230/255., 154/255.],
|
||
'r':[255/255,0,0],
|
||
'g':[0,255/255,0],
|
||
'b':[0,0,255/255],
|
||
'k':[0,0,0],
|
||
'y':[255/255,255/255,0],
|
||
'purple':[128/255,0,128/255]
|
||
}
|
||
|
||
def get_colors(pid):
|
||
if isinstance(pid, int):
|
||
return colors[pid % len(colors)]
|
||
elif isinstance(pid, str):
|
||
return colors_table[pid]
|
||
|
||
from pyrender import RenderFlags
|
||
render_flags = {
|
||
'flip_wireframe': False,
|
||
'all_wireframe': False,
|
||
'all_solid': True,
|
||
'shadows': True,
|
||
'vertex_normals': False,
|
||
'face_normals': False,
|
||
'cull_faces': False, # set to False
|
||
'point_size': 1.0,
|
||
'rgba':True
|
||
}
|
||
|
||
flags = RenderFlags.NONE
|
||
if render_flags['flip_wireframe']:
|
||
flags |= RenderFlags.FLIP_WIREFRAME
|
||
elif render_flags['all_wireframe']:
|
||
flags |= RenderFlags.ALL_WIREFRAME
|
||
elif render_flags['all_solid']:
|
||
flags |= RenderFlags.ALL_SOLID
|
||
|
||
if render_flags['shadows']:
|
||
flags |= RenderFlags.SHADOWS_DIRECTIONAL | RenderFlags.SHADOWS_SPOT
|
||
if render_flags['vertex_normals']:
|
||
flags |= RenderFlags.VERTEX_NORMALS
|
||
if render_flags['face_normals']:
|
||
flags |= RenderFlags.FACE_NORMALS
|
||
if not render_flags['cull_faces']:
|
||
flags |= RenderFlags.SKIP_CULL_FACES
|
||
if render_flags['rgba']:
|
||
flags |= RenderFlags.RGBA
|
||
|
||
|
||
class Renderer(object):
|
||
def __init__(self, focal_length=1000, height=512, width=512, faces=None,
|
||
bg_color=[0.0, 0.0, 0.0, 0.0] # render 配置
|
||
):
|
||
self.renderer = pyrender.OffscreenRenderer(height, width)
|
||
self.faces = faces
|
||
self.focal_length = focal_length
|
||
self.bg_color = bg_color
|
||
self.ambient_light = (0.3, 0.3, 0.3)
|
||
|
||
def add_light(self, scene):
|
||
trans = [0, 0, 0]
|
||
# Use 3 directional lights
|
||
# Create light source
|
||
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
|
||
light_pose = np.eye(4)
|
||
light_pose[:3, 3] = np.array([0, -1, 1]) + trans
|
||
scene.add(light, pose=light_pose)
|
||
light_pose[:3, 3] = np.array([0, 1, 1]) + trans
|
||
scene.add(light, pose=light_pose)
|
||
light_pose[:3, 3] = np.array([1, 1, 2]) + trans
|
||
scene.add(light, pose=light_pose)
|
||
|
||
def render(self, render_data, cameras, images,
|
||
use_white=False,
|
||
ret_depth=False, ret_color=False):
|
||
# Need to flip x-axis
|
||
rot = trimesh.transformations.rotation_matrix(
|
||
np.radians(180), [1, 0, 0])
|
||
output_images, output_colors, output_depths = [], [], []
|
||
for nv, img_ in enumerate(images):
|
||
if use_white:
|
||
img = np.zeros_like(img_, dtype=np.uint8) + 255
|
||
else:
|
||
img = img_.copy()
|
||
K, R, T = cameras['K'][nv], cameras['R'][nv], cameras['T'][nv]
|
||
self.renderer.viewport_height = img.shape[0]
|
||
self.renderer.viewport_width = img.shape[1]
|
||
scene = pyrender.Scene(bg_color=self.bg_color,
|
||
ambient_light=self.ambient_light)
|
||
for trackId, data in render_data.items():
|
||
vert = data['vertices'].copy()
|
||
faces = data['faces']
|
||
# 如果使用了vid这个键，那么可视化的颜色使用vid的颜色
|
||
col = get_colors(data.get('vid', trackId))
|
||
vert = vert @ R.T + T.T
|
||
mesh = trimesh.Trimesh(vert, faces)
|
||
mesh.apply_transform(rot)
|
||
|
||
material = pyrender.MetallicRoughnessMaterial(
|
||
metallicFactor=0.0,
|
||
alphaMode='OPAQUE',
|
||
baseColorFactor=col)
|
||
mesh = pyrender.Mesh.from_trimesh(
|
||
mesh,
|
||
material=material)
|
||
scene.add(mesh, 'mesh')
|
||
camera_pose = np.eye(4)
|
||
camera = pyrender.camera.IntrinsicsCamera(fx=K[0, 0], fy=K[1, 1], cx=K[0, 2], cy=K[1, 2])
|
||
scene.add(camera, pose=camera_pose)
|
||
self.add_light(scene)
|
||
# Alpha channel was not working previously need to check again
|
||
# Until this is fixed use hack with depth image to get the opacity
|
||
rend_rgba, rend_depth = self.renderer.render(scene, flags=flags)
|
||
if rend_rgba.shape[2] == 3: # fail to generate transparent channel
|
||
valid_mask = (rend_depth > 0)[:, :, None]
|
||
rend_rgba = np.dstack((rend_rgba, (valid_mask*255).astype(np.uint8)))
|
||
rend_cat = cv2.addWeighted(cv2.bitwise_and(img, 255 - rend_rgba[:, :, 3:4].repeat(3, 2)), 1, rend_rgba[:, :, :3], 1, 0)
|
||
|
||
output_colors.append(rend_rgba)
|
||
output_depths.append(rend_depth)
|
||
output_images.append(rend_cat)
|
||
if ret_depth:
|
||
return output_images, output_depths
|
||
elif ret_color:
|
||
return output_colors
|
||
else:
|
||
return output_images
|
||
|
||
def render_multiview(self, vertices, K, R, T, imglist, trackId=0, return_depth=False, return_color=False,
|
||
bg_color=[0.0, 0.0, 0.0, 0.0], camera=None):
|
||
# List to store rendered scenes
|
||
output_images, output_colors, output_depths = [], [], []
|
||
# Need to flip x-axis
|
||
rot = trimesh.transformations.rotation_matrix(
|
||
np.radians(180), [1, 0, 0])
|
||
nViews = len(imglist)
|
||
for nv in range(nViews):
|
||
img = imglist[nv]
|
||
self.renderer.viewport_height = img.shape[0]
|
||
self.renderer.viewport_width = img.shape[1]
|
||
# Create a scene for each image and render all meshes
|
||
scene = pyrender.Scene(bg_color=bg_color,
|
||
ambient_light=(0.3, 0.3, 0.3))
|
||
camera_pose = np.eye(4)
|
||
|
||
# for every person in the scene
|
||
if isinstance(vertices, dict):
|
||
for trackId, data in vertices.items():
|
||
vert = data['vertices'].copy()
|
||
faces = data['faces']
|
||
col = data.get('col', trackId)
|
||
vert = vert @ R[nv].T + T[nv]
|
||
mesh = trimesh.Trimesh(vert, faces)
|
||
mesh.apply_transform(rot)
|
||
trans = [0, 0, 0]
|
||
|
||
material = pyrender.MetallicRoughnessMaterial(
|
||
metallicFactor=0.0,
|
||
alphaMode='OPAQUE',
|
||
baseColorFactor=colors[col % len(colors)])
|
||
mesh = pyrender.Mesh.from_trimesh(
|
||
mesh,
|
||
material=material)
|
||
scene.add(mesh, 'mesh')
|
||
else:
|
||
verts = vertices @ R[nv].T + T[nv]
|
||
mesh = trimesh.Trimesh(verts, self.faces)
|
||
mesh.apply_transform(rot)
|
||
trans = [0, 0, 0]
|
||
|
||
material = pyrender.MetallicRoughnessMaterial(
|
||
metallicFactor=0.0,
|
||
alphaMode='OPAQUE',
|
||
baseColorFactor=colors[trackId % len(colors)])
|
||
mesh = pyrender.Mesh.from_trimesh(
|
||
mesh,
|
||
material=material)
|
||
scene.add(mesh, 'mesh')
|
||
|
||
if camera is not None:
|
||
light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=70)
|
||
light_pose = np.eye(4)
|
||
light_pose[:3, 3] = [0, 0, 4.5]
|
||
scene.add(light, pose=light_pose)
|
||
|
||
light_pose[:3, 3] = [0, 1, 4]
|
||
scene.add(light, pose=light_pose)
|
||
|
||
light_pose[:3, 3] = [0, -1, 4]
|
||
scene.add(light, pose=light_pose)
|
||
else:
|
||
trans = [0, 0, 0]
|
||
# Use 3 directional lights
|
||
# Create light source
|
||
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
|
||
light_pose = np.eye(4)
|
||
light_pose[:3, 3] = np.array([0, -1, 1]) + trans
|
||
scene.add(light, pose=light_pose)
|
||
light_pose[:3, 3] = np.array([0, 1, 1]) + trans
|
||
scene.add(light, pose=light_pose)
|
||
light_pose[:3, 3] = np.array([1, 1, 2]) + trans
|
||
scene.add(light, pose=light_pose)
|
||
if camera is None:
|
||
if K is None:
|
||
camera_center = np.array([img.shape[1] / 2., img.shape[0] / 2.])
|
||
camera = pyrender.camera.IntrinsicsCamera(fx=self.focal_length, fy=self.focal_length, cx=camera_center[0], cy=camera_center[1])
|
||
else:
|
||
camera = pyrender.camera.IntrinsicsCamera(fx=K[nv][0, 0], fy=K[nv][1, 1], cx=K[nv][0, 2], cy=K[nv][1, 2])
|
||
scene.add(camera, pose=camera_pose)
|
||
# Alpha channel was not working previously need to check again
|
||
# Until this is fixed use hack with depth image to get the opacity
|
||
color, rend_depth = self.renderer.render(scene, flags=flags)
|
||
# color = color[::-1,::-1]
|
||
# rend_depth = rend_depth[::-1,::-1]
|
||
output_depths.append(rend_depth)
|
||
color = color.astype(np.uint8)
|
||
valid_mask = (rend_depth > 0)[:, :, None]
|
||
if color.shape[2] == 3: # 在服务器上透明通道失败
|
||
color = np.dstack((color, (valid_mask*255).astype(np.uint8)))
|
||
output_colors.append(color)
|
||
output_img = (color[:, :, :3] * valid_mask +
|
||
(1 - valid_mask) * img)
|
||
|
||
output_img = output_img.astype(np.uint8)
|
||
output_images.append(output_img)
|
||
if return_depth:
|
||
return output_images, output_depths
|
||
elif return_color:
|
||
return output_colors
|
||
else:
|
||
return output_images
|
||
|
||
def render_results(img, render_data, cam_params, outname=None, rotate=False, degree=90, axis=[1.,0.,0],
|
||
fix_center=None):
|
||
render_data = copy.deepcopy(render_data)
|
||
render = Renderer(height=1024, width=1024, faces=None)
|
||
Ks, Rs, Ts = [cam_params['K']], [cam_params['Rc']], [cam_params['Tc']]
|
||
imgsrender = render.render_multiview(render_data, Ks, Rs, Ts, [img], return_color=True)[0]
|
||
render0 = cv2.addWeighted(cv2.bitwise_and(img, 255 - imgsrender[:, :, 3:4].repeat(3, 2)), 1, imgsrender[:, :, :3], 1, 0.0)
|
||
if rotate:
|
||
# simple rotate the vertices
|
||
if fix_center is None:
|
||
center = np.mean(np.vstack([v['vertices'] for i, v in render_data.items()]), axis=0, keepdims=True)
|
||
new_center = center.copy()
|
||
new_center[:, 0:2] = 0
|
||
else:
|
||
center = fix_center.copy()
|
||
new_center = fix_center.copy()
|
||
new_center[:, 2] *= 1.5
|
||
direc = np.array(axis)
|
||
rot, _ = cv2.Rodrigues(direc*degree/90*np.pi/2)
|
||
for key in render_data.keys():
|
||
vertices = render_data[key]['vertices']
|
||
vert = (vertices - center) @ rot.T + new_center
|
||
render_data[key]['vertices'] = vert
|
||
blank = np.zeros(())
|
||
blank = np.zeros((img.shape[0], img.shape[1], 3), dtype=img.dtype) + 255
|
||
imgsrender = render.render_multiview(render_data, Ks, Rs, Ts, [blank], return_color=True)[0]
|
||
render1 = cv2.addWeighted(cv2.bitwise_and(blank, 255- imgsrender[:, :, 3:4].repeat(3, 2)), 1, imgsrender[:, :, :3], 1, 0.0)
|
||
render0 = np.vstack([render0, render1])
|
||
if outname is not None:
|
||
os.makedirs(os.path.dirname(outname), exist_ok=True)
|
||
cv2.imwrite(outname, render0)
|
||
return render0 |