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': False, # TODO:bug exists in shadow mode '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=[1.0, 1.0, 1.0, 0.0], down_scale=1, # render 配置 extra_mesh=[] ): self.renderer = pyrender.OffscreenRenderer(height, width) self.faces = faces self.focal_length = focal_length self.bg_color = bg_color self.ambient_light = (0.5, 0.5, 0.5) self.down_scale = down_scale self.extra_mesh = extra_mesh 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=3) light_forward = np.eye(4) light_forward[:3, 3] = np.array([0, 1, 1]) scene.add(light, pose=light_forward) light_z = np.eye(4) light_z[:3, :3] = cv2.Rodrigues(np.array([-np.pi/2, 0, 0]))[0] scene.add(light, pose=light_z) def render(self, render_data, cameras, images, use_white=False, add_back=True, 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].copy(), cameras['R'][nv], cameras['T'][nv] # down scale the image to speed up rendering img = cv2.resize(img, None, fx=1/self.down_scale, fy=1/self.down_scale) K[:2, :] /= self.down_scale 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 iextra, _mesh in enumerate(self.extra_mesh): if True: mesh = _mesh.copy() trans_cam = np.eye(4) trans_cam[:3, :3] = R trans_cam[:3, 3:] = T mesh.apply_transform(trans_cam) mesh.apply_transform(rot) # mesh.vertices = np.asarray(mesh.vertices) @ R.T + T.T mesh_ = pyrender.Mesh.from_trimesh(mesh) scene.add(mesh_, 'extra{}'.format(iextra)) else: vert = np.asarray(_mesh.vertices).copy() faces = np.asarray(_mesh.faces) vert = vert @ R.T + T.T mesh = trimesh.Trimesh(vert, faces, process=False) mesh.apply_transform(rot) material = pyrender.MetallicRoughnessMaterial( metallicFactor=0.0, alphaMode='OPAQUE', baseColorFactor=(0., 0., 0., 1.)) mesh = pyrender.Mesh.from_trimesh( mesh, material=material) scene.add(mesh, 'extra{}'.format(iextra)) for trackId, data in render_data.items(): vert = data['vertices'].copy() faces = data['faces'] vert = vert @ R.T + T.T if 'colors' not in data.keys(): # 如果使用了vid这个键，那么可视化的颜色使用vid的颜色 col = get_colors(data.get('vid', trackId)) mesh = trimesh.Trimesh(vert, faces, process=False) 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, data['name']) else: mesh = trimesh.Trimesh(vert, faces, vertex_colors=data['colors'], process=False) mesh.apply_transform(rot) material = pyrender.MetallicRoughnessMaterial( metallicFactor=0.0, alphaMode='OPAQUE', baseColorFactor=(1., 1., 1.)) mesh = pyrender.Mesh.from_trimesh(mesh, material=material) scene.add(mesh, data['name']) 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) # pyrender.Viewer(scene, use_raymond_lighting=True) # 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_rgba = rend_rgba[..., [2, 1, 0, 3]] if add_back: rend_cat = cv2.addWeighted( cv2.bitwise_and(img, 255 - rend_rgba[:, :, 3:4].repeat(3, 2)), 1, cv2.bitwise_and(rend_rgba[:, :, :3], rend_rgba[:, :, 3:4].repeat(3, 2)), 1, 0) else: rend_cat = rend_rgba 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