Spaces:
Sleeping
Sleeping
File size: 5,060 Bytes
99a05f0 |
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 |
# 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()
|