-
Notifications
You must be signed in to change notification settings - Fork 0
/
ModelPartsScreenshot.py
82 lines (62 loc) · 2.56 KB
/
ModelPartsScreenshot.py
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
import copy
import cv2
import numpy as np
import ProjectUtils as pUtils
import pyrender
from pyrender.constants import *
# Permanently changes the rotation
def rotateAllMeshesInScene(scene, rotationRad):
for meshNode in scene.mesh_nodes:
quaternionRotation = pUtils.radEuler2Quat(rotationRad)
meshNode.rotation = quaternionRotation
# Rotates meshes, takes picture and resets the rotation back
def renderSceneWithRotation(offscreenRenderer, scene, rotationRad, depthBegin, depthEnd):
rotateAllMeshesInScene(scene, rotationRad)
#flags = RenderFlags.DEPTH_ONLY
pixels, depth = offscreenRenderer.render(scene)
rows = depth.shape[0]
cols = depth.shape[1]
for x in range(0, rows):
for y in range(0, cols):
depthValue = depth[x][y]
# make it white if no depth is present
if depthValue == 0.0:
depth[x][y] = 255
else:
# at begin depth color would be the darkest, at end - the lightest
value = pUtils.inverseLerp(depthBegin, depthEnd, depthValue)
depth[x][y] = 255 * value
# convert to int
depth = depth.astype(np.uint8)
# turn it to w,h,3 shape
#depth = cv2.cvtColor(depth, cv2.COLOR_GRAY2RGB)
#depth = cv2.cvtColor(depth, cv2.COLOR_RGB2GRAY)
rotateAllMeshesInScene(scene, (0, 0, 0))
return depth
def captureDepth(model, rotations, imageWidth=224, imageHeight=224, cameraZTranslation=2.5, lightIntensity=2.0, depthBegin=1, depthEnd=5):
# Construct an offline scene
scene = pyrender.Scene()
# add parts
for part in model.parts:
partMesh = pyrender.Mesh.from_trimesh(part.mesh, smooth=False)
scene.add(partMesh)
# add camera
renderCamera = pyrender.PerspectiveCamera(
yfov=np.pi / 3.0, aspectRatio=imageWidth/imageHeight)
cameraNode = pyrender.Node(camera=renderCamera, matrix=np.eye(4))
cameraNode.translation[2] = cameraZTranslation
scene.add_node(cameraNode)
# add light
light = pyrender.DirectionalLight(
color=[1.0, 1.0, 1.0], intensity=lightIntensity)
lightNode = pyrender.Node(light=light, matrix=np.eye(4))
scene.add_node(lightNode)
# initialize offscreen renderer
offscreenRenderer = pyrender.OffscreenRenderer(
viewport_width=imageWidth, viewport_height=imageHeight, point_size=1.0)
# render
result = []
for rotation in rotations:
result.append(renderSceneWithRotation(offscreenRenderer,
scene, rotation, depthBegin, depthEnd))
return result