-
Notifications
You must be signed in to change notification settings - Fork 0
/
generation.py
383 lines (312 loc) · 13.1 KB
/
generation.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
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
import torch
import torch.optim as optim
from torch import autograd
import numpy as np
from tqdm import trange
import trimesh
from utils import libmcubes
from common import make_3d_grid
from utils.libsimplify import simplify_mesh
from utils.libmise import MISE
import time
import math
class Generator3D(object):
''' Generator class for Occupancy Networks.
It provides functions to generate the final mesh as well refining options.
Args:
model (nn.Module): trained Occupancy Network model
points_batch_size (int): batch size for points evaluation
threshold (float): threshold value
refinement_step (int): number of refinement steps
resolution0 (int): start resolution for MISE
upsampling steps (int): number of upsampling steps
with_normals (bool): whether normals should be estimated
padding (float): how much padding should be used for MISE
simplify_nfaces (int): number of faces the mesh should be simplified to
preprocessor (nn.Module): preprocessor for inputs
'''
def __init__(self, model, points_batch_size=100000,
threshold=0.5, refinement_step=0,
resolution0=16, upsampling_steps=3,
with_normals=False, padding=0.1,
simplify_nfaces=None,
preprocessor=None):
self.model = model.cuda()
self.points_batch_size = points_batch_size
self.refinement_step = refinement_step
self.threshold = threshold
self.resolution0 = resolution0
self.upsampling_steps = upsampling_steps
self.with_normals = with_normals
self.padding = padding
self.simplify_nfaces = simplify_nfaces
self.preprocessor = preprocessor
def generate_mesh(self, inputs, return_stats=True):
''' Generates the output mesh.
Args:
data (tensor): data tensor
return_stats (bool): whether stats should be returned
'''
self.model.eval()
stats_dict = {}
kwargs = {}
# Preprocess if requires
if self.preprocessor is not None:
t0 = time.time()
with torch.no_grad():
inputs = self.preprocessor(inputs)
stats_dict['time (preprocess)'] = time.time() - t0
# Encode inputs
t0 = time.time()
with torch.no_grad():
c = self.model.encode_inputs(inputs)
stats_dict['time (encode inputs)'] = time.time() - t0
mesh = self.generate_from_latent(c, stats_dict=stats_dict, **kwargs)
if return_stats:
return mesh, stats_dict
else:
return mesh
def generate_voxel(self, inputs):
''' Generates the output mesh.
Args:
data (tensor): data tensor
return_stats (bool): whether stats should be returned
'''
# model.evel() before using this function: generate_voxel()
self.model.eval()
# stats_dict = {}
kwargs = {}
# Preprocess if requires
if self.preprocessor is not None:
# t0 = time.time()
with torch.no_grad():
inputs = self.preprocessor(inputs)
# stats_dict['time (preprocess)'] = time.time() - t0
# Encode inputs
# t0 = time.time()
with torch.no_grad():
c = self.model.encode_inputs(inputs)
# stats_dict['time (encode inputs)'] = time.time() - t0
# Compute bounding box size
box_size = 1
if self.upsampling_steps == 0:
nx = self.resolution0
pointsf = box_size * make_3d_grid(
(-0.5,) * 3, (0.5,) * 3, (nx,) * 3
)
values = self.eval_points(pointsf, c, **kwargs).cpu().numpy()
values = 1 / (1 + np.exp(-values))
value_grid = values.reshape(nx, nx, nx)
else:
mesh_extractor = MISE(self.resolution0, self.upsampling_steps, self.threshold)
points = mesh_extractor.query()
while points.shape[0] != 0:
# Query points
pointsf = torch.FloatTensor(points).cuda()
# Normalize to bounding box
pointsf = pointsf / mesh_extractor.resolution
pointsf = box_size * (pointsf - 0.5)
# Evaluate model and update
values = self.eval_points(pointsf, c, **kwargs).cpu().numpy()
values = 1 / (1 + np.exp(-values))
values = values.astype(np.float64)
mesh_extractor.update(points, values)
points = mesh_extractor.query()
value_grid = mesh_extractor.to_dense()
return value_grid
def generate_from_latent(self, c=None, stats_dict={}, **kwargs):
''' Generates mesh from latent.
Args:
c (tensor): latent conditioned code c
stats_dict (dict): stats dictionary
'''
threshold = np.log(self.threshold) - np.log(1. - self.threshold)
t0 = time.time()
# Compute bounding box size
box_size = 1 + self.padding
# Shortcut
if self.upsampling_steps == 0:
nx = self.resolution0
pointsf = box_size * make_3d_grid(
(-0.5,)*3, (0.5,)*3, (nx,)*3
)
values = self.eval_points(pointsf, c, **kwargs).cpu().numpy()
value_grid = values.reshape(nx, nx, nx)
else:
mesh_extractor = MISE(self.resolution0, self.upsampling_steps, threshold)
points = mesh_extractor.query()
while points.shape[0] != 0:
# Query points
pointsf = torch.FloatTensor(points).cuda()
# Normalize to bounding box
pointsf = pointsf / mesh_extractor.resolution
pointsf = box_size * (pointsf - 0.5)
# Evaluate model and update
values = self.eval_points(pointsf, c, **kwargs).cpu().numpy()
values = values.astype(np.float64)
mesh_extractor.update(points, values)
points = mesh_extractor.query()
value_grid = mesh_extractor.to_dense()
# Extract mesh
stats_dict['time (eval points)'] = time.time() - t0
mesh = self.extract_mesh(value_grid, c, stats_dict=stats_dict)
return mesh
def eval_points(self, p, c=None, **kwargs):
''' Evaluates the occupancy values for the points.
Args:
p (tensor): points
c (tensor): latent conditioned code c
'''
p_split = torch.split(p, self.points_batch_size)
# print('p_shape', p.size())
occ_hats = []
for pi in p_split:
pi = pi.unsqueeze(0).cuda()
with torch.no_grad():
occ_hat = self.model.decode(pi, c, **kwargs).logits
occ_hats.append(occ_hat.squeeze(0).detach().cpu())
occ_hat = torch.cat(occ_hats, dim=0)
# print('occ_shape:', occ_hat.size())
return occ_hat
def extract_mesh(self, occ_hat, c=None, stats_dict=dict()):
''' Extracts the mesh from the predicted occupancy grid.
Args:
occ_hat (tensor): value grid of occupancies
c (tensor): latent conditioned code c
stats_dict (dict): stats dictionary
'''
# Some short hands
n_x, n_y, n_z = occ_hat.shape
box_size = 1 + self.padding
threshold = np.log(self.threshold) - np.log(1. - self.threshold)
# ### just for testing ###
# # print('occ_val:', occ_hat)
# occ_hat_sigmoid = 1./(1.+np.exp(-1.*occ_hat))
# # print('occ_val_sigmoid:', occ_hat_sigmoid)
# print('occ_shape:', occ_hat_sigmoid.shape)
# threshold_sigmoid = 1./(1.+np.exp(-threshold))
# print('threshold:', threshold, threshold_sigmoid)
# print('occ_val_sigmoid_min_max:', occ_hat_sigmoid.min(), occ_hat_sigmoid.max())
# print('occ_val_sigmoid_average:', occ_hat_sigmoid.mean())
# with open('occ_val.obj', 'w') as f:
# for i in range(0, occ_hat_sigmoid.shape[0]):
# for j in range(0, occ_hat_sigmoid.shape[1]):
# for k in range(0, occ_hat_sigmoid.shape[2]):
# if (occ_hat_sigmoid[i][j][k] > threshold_sigmoid):
# f.write('v ' + str(i) + ' ' + str(j) + ' ' + str(k) + '\n')
# f.close()
#
# # print('occ_val_sigmoid:', occ_hat_sigmoid)
# ### just for testing ###
# Make sure that mesh is watertight
t0 = time.time()
occ_hat_padded = np.pad(
occ_hat, 1, 'constant', constant_values=-1e6)
vertices, triangles = libmcubes.marching_cubes(
occ_hat_padded, threshold)
stats_dict['time (marching cubes)'] = time.time() - t0
# Strange behaviour in libmcubes: vertices are shifted by 0.5
vertices -= 0.5
# Undo padding
vertices -= 1
# Normalize to bounding box
vertices /= np.array([n_x-1, n_y-1, n_z-1])
vertices = box_size * (vertices - 0.5)
# mesh_pymesh = pymesh.form_mesh(vertices, triangles)
# mesh_pymesh = fix_pymesh(mesh_pymesh)
# Estimate normals if needed
if self.with_normals and not vertices.shape[0] == 0:
t0 = time.time()
normals = self.estimate_normals(vertices, c)
stats_dict['time (normals)'] = time.time() - t0
else:
normals = None
# Create mesh
mesh = trimesh.Trimesh(vertices, triangles,
vertex_normals=normals,
process=False)
# Directly return if mesh is empty
if vertices.shape[0] == 0:
return mesh
# TODO: normals are lost here
if self.simplify_nfaces is not None:
t0 = time.time()
mesh = simplify_mesh(mesh, self.simplify_nfaces, 5.)
stats_dict['time (simplify)'] = time.time() - t0
# Refine mesh
if self.refinement_step > 0:
t0 = time.time()
self.refine_mesh(mesh, occ_hat, c)
stats_dict['time (refine)'] = time.time() - t0
return mesh
def estimate_normals(self, vertices, c=None):
''' Estimates the normals by computing the gradient of the objective.
Args:
vertices (numpy array): vertices of the mesh
c (tensor): latent conditioned code c
'''
vertices = torch.FloatTensor(vertices)
vertices_split = torch.split(vertices, self.points_batch_size)
normals = []
c = c.unsqueeze(0)
for vi in vertices_split:
vi = vi.unsqueeze(0).cuda()
vi.requires_grad_()
occ_hat = self.model.decode(vi, c).logits
out = occ_hat.sum()
out.backward()
ni = -vi.grad
ni = ni / torch.norm(ni, dim=-1, keepdim=True)
ni = ni.squeeze(0).cpu().numpy()
normals.append(ni)
normals = np.concatenate(normals, axis=0)
return normals
def refine_mesh(self, mesh, occ_hat, c=None):
''' Refines the predicted mesh.
Args:
mesh (trimesh object): predicted mesh
occ_hat (tensor): predicted occupancy grid
c (tensor): latent conditioned code c
'''
self.model.eval()
# Some shorthands
n_x, n_y, n_z = occ_hat.shape
assert(n_x == n_y == n_z)
# threshold = np.log(self.threshold) - np.log(1. - self.threshold)
threshold = self.threshold
# Vertex parameter
v0 = torch.FloatTensor(mesh.vertices).cuda()
v = torch.nn.Parameter(v0.clone())
# Faces of mesh
faces = torch.LongTensor(mesh.faces).cuda()
# Start optimization
optimizer = optim.RMSprop([v], lr=1e-4)
for it_r in trange(self.refinement_step):
optimizer.zero_grad()
# Loss
face_vertex = v[faces]
eps = np.random.dirichlet((0.5, 0.5, 0.5), size=faces.shape[0])
eps = torch.FloatTensor(eps).cuda()
face_point = (face_vertex * eps[:, :, None]).sum(dim=1)
face_v1 = face_vertex[:, 1, :] - face_vertex[:, 0, :]
face_v2 = face_vertex[:, 2, :] - face_vertex[:, 1, :]
face_normal = torch.cross(face_v1, face_v2)
face_normal = face_normal / \
(face_normal.norm(dim=1, keepdim=True) + 1e-10)
face_value = torch.sigmoid(
self.model.decode(face_point.unsqueeze(0), c).logits
)
normal_target = -autograd.grad(
[face_value.sum()], [face_point], create_graph=True)[0]
normal_target = \
normal_target / \
(normal_target.norm(dim=1, keepdim=True) + 1e-10)
loss_target = (face_value - threshold).pow(2).mean()
loss_normal = \
(face_normal - normal_target).pow(2).sum(dim=1).mean()
loss = loss_target + 0.01 * loss_normal
# Update
loss.backward()
optimizer.step()
mesh.vertices = v.data.cpu().numpy()
return mesh