-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain_render_single_view.py
204 lines (166 loc) · 6.52 KB
/
main_render_single_view.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
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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
import numpy as np
import os
import cv2 as cv
import glob
import math
import random
import pyexr
from tqdm import tqdm
import scipy.io as sio
import prt.sh_util as sh_util
import gc
from renderer.camera import Camera
from renderer.mesh import load_obj_mesh, compute_tangent, compute_normal, load_obj_mesh_mtl
"""
runtime configuration
"""
mesh_data_dir = os.path.join(os.path.dirname(__file__), '../../../dataset_example/mesh_data')
output_data_dir = os.path.join(os.path.dirname(__file__), '../../../dataset_example/image_data')
view_num = 180
cam_f = 5000
cam_dist = 10
img_res = 512
def get_data_list():
"""reads data list"""
data_list = glob.glob(os.path.join(mesh_data_dir, './*/'))
return sorted(data_list)
def read_data(item):
"""reads data """
mesh_filename = glob.glob(os.path.join(item, '*.obj'))[0] # assumes one .obj file
text_filename = glob.glob(os.path.join(item, '*.jpg'))[0] # assumes one .jpg file
vertices, faces, normals, faces_normals, textures, face_textures \
= load_obj_mesh(mesh_filename, with_normal=True, with_texture=True)
texture_image = cv.imread(text_filename)
texture_image = cv.cvtColor(texture_image, cv.COLOR_BGR2RGB)
prt_data = sio.loadmat(os.path.join(item, 'bounce/prt_data.mat'))
prt, face_prt = prt_data['bounce0'], prt_data['face']
return vertices, faces, normals, faces_normals, textures, face_textures, texture_image, prt, face_prt
def make_rotate(rx, ry, rz):
sinX = np.sin(rx)
sinY = np.sin(ry)
sinZ = np.sin(rz)
cosX = np.cos(rx)
cosY = np.cos(ry)
cosZ = np.cos(rz)
Rx = np.zeros((3,3))
Rx[0, 0] = 1.0
Rx[1, 1] = cosX
Rx[1, 2] = -sinX
Rx[2, 1] = sinX
Rx[2, 2] = cosX
Ry = np.zeros((3,3))
Ry[0, 0] = cosY
Ry[0, 2] = sinY
Ry[1, 1] = 1.0
Ry[2, 0] = -sinY
Ry[2, 2] = cosY
Rz = np.zeros((3,3))
Rz[0, 0] = cosZ
Rz[0, 1] = -sinZ
Rz[1, 0] = sinZ
Rz[1, 1] = cosZ
Rz[2, 2] = 1.0
R = np.matmul(np.matmul(Rz,Ry),Rx)
return R
def generate_cameras(dist=10, view_num=60):
cams = []
target = [0, 0, 0]
up = [0, 1, 0]
for view_idx in range(view_num):
angle = (math.pi * 2 / view_num) * view_idx
eye = np.asarray([dist * math.sin(angle), 0, dist * math.cos(angle)])
fwd = np.asarray(target, np.float64) - eye
fwd /= np.linalg.norm(fwd)
right = np.cross(fwd, up)
right /= np.linalg.norm(right)
down = np.cross(fwd, right)
cams.append(
{
'center': eye,
'direction': fwd,
'right': right,
'up': -down,
}
)
return cams
def process_one_data_item(data_item, rndr, rndr_uv, shs, view_idx):
_, item_name = os.path.split(data_item[:-1])
output_fd = os.path.join(output_data_dir, item_name)
os.makedirs(output_fd, exist_ok=True)
os.makedirs(os.path.join(output_fd, 'color'), exist_ok=True)
os.makedirs(os.path.join(output_fd, 'mask'), exist_ok=True)
# os.makedirs(os.path.join(output_fd, 'color_uv'), exist_ok=True)
os.makedirs(os.path.join(output_fd, 'meta'), exist_ok=True)
vertices, faces, normals, faces_normals, textures, face_textures, \
texture_image, prt, face_prt = read_data(data_item)
cam = Camera(width=img_res, height=img_res, focal=5000, near=0.1, far=40)
cam.sanity_check()
rndr.set_norm_mat(1.0, 0.0)
tan, bitan = compute_tangent(vertices, faces, normals, textures, face_textures)
rndr.set_mesh(vertices, faces, normals, faces_normals, textures, face_textures, prt, face_prt, tan, bitan)
rndr.set_albedo(texture_image)
# rndr_uv.set_mesh(vertices, faces, normals, faces_normals, textures, face_textures, prt, face_prt, tan, bitan)
# rndr_uv.set_albedo(texture_image)
cam_params = generate_cameras(dist=cam_dist, view_num=view_num)
sh_list = []
cam_param = cam_params[view_idx]
cam.center = cam_param['center']
cam.right = cam_param['right']
cam.up = cam_param['up']
cam.direction = cam_param['direction']
cam.sanity_check()
rndr.set_camera(cam)
rndr_uv.set_camera(cam)
sh_id = random.randint(0,shs.shape[0]-1)
sh = shs[sh_id]
sh_angle = 0.2*np.pi*(random.random()-0.5)
sh = sh_util.rotateSH(sh, make_rotate(0, sh_angle, 0).T)
sh_list.append(sh)
rndr.set_sh(sh)
rndr.analytic = False
rndr.use_inverse_depth = False
rndr.display()
out_all_f = rndr.get_color(0)
out_mask = out_all_f[:,:,3]
out_all_f = cv.cvtColor(out_all_f, cv.COLOR_RGBA2BGR)
cv.imwrite(os.path.join(output_fd, 'color', '%04d.jpg' % view_idx), np.uint8(out_all_f * 255))
cv.imwrite(os.path.join(output_fd, 'mask', '%04d.png' % view_idx), np.uint8(out_mask * 255))
# rndr_uv.set_sh(sh)
# rndr_uv.analytic = False
# rndr_uv.use_inverse_depth = False
# rndr_uv.display()
# uv_color = rndr_uv.get_color(0)
# uv_color = cv.cvtColor(uv_color, cv.COLOR_RGBA2BGR)
# cv.imwrite(os.path.join(output_fd, 'color_uv', '%04d.png' % view_idx), np.uint8(uv_color * 255))
if view_idx == 0:
uv_pos = rndr_uv.get_color(1)
uv_mask = uv_pos[:,:,3]
cv.imwrite(os.path.join(output_fd, 'meta', 'uv_mask.png'), np.uint8(uv_mask * 255))
data = {'default': uv_pos[:,:,:3]} # default is a reserved name
pyexr.write(os.path.join(output_fd, 'meta', 'uv_pos.exr'), data)
uv_nml = rndr_uv.get_color(2)
uv_nml = cv.cvtColor(uv_nml, cv.COLOR_RGBA2BGR)
cv.imwrite(os.path.join(output_fd, 'meta', 'uv_nml.png'), np.uint8(uv_nml * 255))
sio.savemat(
os.path.join(output_fd, 'meta', 'cam_data.mat'),
{'cam': cam_params})
sio.savemat(
os.path.join(output_fd, 'meta', 'sh_data.mat'),
{'sh': sh_list})
# Clean up resources after rendering each view by manually releasing resources and triggering garbage collection
del rndr
gc.collect()
def main(view_idx):
shs = np.load(os.path.join(os.path.dirname(__file__), 'env_sh.npy'))
egl = False
from renderer.gl.init_gl import initialize_GL_context
initialize_GL_context(width=img_res, height=img_res, egl=egl)
from renderer.gl.prt_render import PRTRender
rndr = PRTRender(width=img_res, height=img_res, ms_rate=1.0, egl=egl)
rndr_uv = PRTRender(width=img_res, height=img_res, uv_mode=True, egl=egl)
data_list = get_data_list()
for data_item in tqdm(data_list, ascii=True):
process_one_data_item(data_item, rndr, rndr_uv, shs, view_idx=view_idx)
print('Done')
if __name__ == '__main__':
main()