Spaces:
Sleeping
Sleeping
| # Copyright (C) 2022. Huawei Technologies Co., Ltd. All rights reserved. | |
| # This program is free software; you can redistribute it and/or modify it | |
| # under the terms of the MIT license. | |
| # This program is distributed in the hope that it will be useful, but WITHOUT ANY | |
| # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A | |
| # PARTICULAR PURPOSE. See the MIT License for more details. | |
| import os | |
| import trimesh | |
| import pyrender | |
| import numpy as np | |
| import colorsys | |
| import cv2 | |
| class Renderer(object): | |
| def __init__(self, focal_length=600, img_w=512, img_h=512, faces=None, | |
| same_mesh_color=False): | |
| os.environ['PYOPENGL_PLATFORM'] = 'egl' | |
| self.renderer = pyrender.OffscreenRenderer(viewport_width=img_w, | |
| viewport_height=img_h, | |
| point_size=1.0) | |
| self.camera_center = [img_w // 2, img_h // 2] | |
| self.focal_length = focal_length | |
| self.faces = faces | |
| self.same_mesh_color = same_mesh_color | |
| def render_front_view(self, verts, bg_img_rgb=None, bg_color=(0, 0, 0, 0), vertex_colors=None, render_part_seg=False, part_label_bins=None): | |
| # Create a scene for each image and render all meshes | |
| scene = pyrender.Scene(bg_color=bg_color, ambient_light=np.ones(3) * (1 if render_part_seg else 0)) | |
| # Create camera. Camera will always be at [0,0,0] | |
| camera = pyrender.camera.IntrinsicsCamera(fx=self.focal_length, fy=self.focal_length, | |
| cx=self.camera_center[0], cy=self.camera_center[1]) | |
| scene.add(camera, pose=np.eye(4)) | |
| # Create light source | |
| if not render_part_seg: | |
| light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=3.0) | |
| # for DirectionalLight, only rotation matters | |
| light_pose = trimesh.transformations.rotation_matrix(np.radians(-45), [1, 0, 0]) | |
| scene.add(light, pose=light_pose) | |
| light_pose = trimesh.transformations.rotation_matrix(np.radians(45), [0, 1, 0]) | |
| scene.add(light, pose=light_pose) | |
| # Need to flip x-axis | |
| rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0]) | |
| # multiple person | |
| num_people = len(verts) | |
| # for every person in the scene | |
| for n in range(num_people): | |
| mesh = trimesh.Trimesh(verts[n], self.faces, process=False) | |
| mesh.apply_transform(rot) | |
| if self.same_mesh_color: | |
| mesh_color = colorsys.hsv_to_rgb(0.6, 0.5, 1.0) | |
| else: | |
| mesh_color = colorsys.hsv_to_rgb(float(n) / num_people, 0.5, 1.0) | |
| material = pyrender.MetallicRoughnessMaterial( | |
| metallicFactor=1.0, | |
| alphaMode='OPAQUE', | |
| baseColorFactor=mesh_color) | |
| if vertex_colors is not None: | |
| # color individual vertices based on part labels | |
| mesh.visual.vertex_colors = vertex_colors | |
| mesh = pyrender.Mesh.from_trimesh(mesh, material=material, wireframe=False) | |
| scene.add(mesh, 'mesh') | |
| # Alpha channel was not working previously, need to check again | |
| # Until this is fixed use hack with depth image to get the opacity | |
| color_rgba, depth_map = self.renderer.render(scene, flags=pyrender.RenderFlags.RGBA) | |
| color_rgb = color_rgba[:, :, :3] | |
| if render_part_seg: | |
| body_parts = color_rgb.copy() | |
| # make single channel | |
| body_parts = body_parts.max(-1) # reduce to single channel | |
| # convert pixel value to bucket indices | |
| # body_parts = torch.bucketize(body_parts, self.part_label_bins, right=True) | |
| body_parts = np.digitize(body_parts, part_label_bins, right=True) | |
| # part labels start from 2 because of the binning scheme. Subtract 1 from all non-zero labels to make label | |
| # go from 1 to 24. 0 is background | |
| # handle background coinciding with hip label = 0 | |
| body_parts = body_parts + 1 | |
| mask = depth_map > 0 | |
| body_parts = body_parts * mask | |
| return body_parts, color_rgb | |
| if bg_img_rgb is None: | |
| return color_rgb | |
| else: | |
| mask = depth_map > 0 | |
| bg_img_rgb[mask] = color_rgb[mask] | |
| return bg_img_rgb | |
| def render_side_view(self, verts): | |
| centroid = verts.mean(axis=(0, 1)) # n*6890*3 -> 3 | |
| # make the centroid at the image center (the X and Y coordinates are zeros) | |
| centroid[:2] = 0 | |
| aroundy = cv2.Rodrigues(np.array([0, np.radians(90.), 0]))[0][np.newaxis, ...] # 1*3*3 | |
| pred_vert_arr_side = np.matmul((verts - centroid), aroundy) + centroid | |
| side_view = self.render_front_view(pred_vert_arr_side) | |
| return side_view | |
| def delete(self): | |
| """ | |
| Need to delete before creating the renderer next time | |
| """ | |
| self.renderer.delete() | |