-
Notifications
You must be signed in to change notification settings - Fork 1
/
utils.py
478 lines (381 loc) · 14.4 KB
/
utils.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
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
import torch
import numpy as np
from laspy.file import File
import plotly.graph_objects as go
import math
import laspy
import torch
import yaml
import dash
import dash_core_components as dcc
import dash_html_components as html
import open3d as o3d
from knn import get_knn
# Losses from original repo
eps = 1e-8
def load_las(path, extra_dim_list=None, scale_colors=True):
"""Load las/laz from given path, laz fiels require laszip on path"""
input_las = File(path, mode='r')
point_records = input_las.points.copy()
las_scaleX = input_las.header.scale[0]
las_offsetX = input_las.header.offset[0]
las_scaleY = input_las.header.scale[1]
las_offsetY = input_las.header.offset[1]
las_scaleZ = input_las.header.scale[2]
las_offsetZ = input_las.header.offset[2]
# calculating coordinates
if scale_colors:
color_div = 65536
else:
color_div = 1
p_X = np.array((point_records['point']['X'] * las_scaleX) + las_offsetX)
p_Y = np.array((point_records['point']['Y'] * las_scaleY) + las_offsetY)
p_Z = np.array((point_records['point']['Z'] * las_scaleZ) + las_offsetZ)
try:
points = np.vstack((p_X, p_Y, p_Z, input_las.red/color_div,
input_las.green/color_div, input_las.blue/color_div)).T
except:
pass
return points
def bits_per_dim(log_likelihood, dims_prod):
multiplier = torch.log(torch.Tensor([2])).to(log_likelihood.device)
bpd = -log_likelihood * multiplier / dims_prod
return bpd
def figure_dash(fig):
app = dash.Dash(name='plot_fig', suppress_callback_exceptions=False)
app.layout = html.Div([dcc.Graph(id='fig', figure=fig, style={
'width': '100vw', 'height': '100vh'})], style={'width': '100vw', 'height': '100vh'})
app.run_server(debug=True)
def view_cloud_plotly(points, rgb=None, fig=None, point_size=5, show=True, axes=False, show_scale=False, colorscale=None, title=None):
"""Creat plotly figure of given cloud"""
if isinstance(points, torch.Tensor):
points = points.cpu()
points = points.detach().numpy()
if isinstance(rgb, torch.Tensor):
rgb = rgb.cpu().detach().numpy()
if rgb is None:
rgb = np.zeros_like(points)
elif rgb.shape[-1] == 1:
rgb = rgb.squeeze()
else:
divide_by = np.maximum(rgb.max(axis=0), eps)
rgb = np.rint(np.divide(rgb, divide_by)*255).astype(np.uint8)
if fig == None:
fig = go.Figure()
if points.shape[1] == 2:
z = np.zeros_like(points[:, 0])
else:
z = points[:, 2]
fig.add_scatter3d(
x=points[:, 0],
y=points[:, 1],
z=z,
marker=dict(
size=point_size,
color=rgb,
colorscale=colorscale,
showscale=show_scale,
opacity=1,
),
opacity=1,
mode='markers',
)
if not axes:
fig.update_layout(
scene=dict(
xaxis=dict(showticklabels=False, visible=False),
yaxis=dict(showticklabels=False, visible=False),
zaxis=dict(showticklabels=False, visible=False)
))
if title != None:
fig.update_layout(title_text=title)
if show:
figure_dash(fig)
# fig.show()
return fig
def extract_area(full_cloud, center, clearance, shape='circle'):
"""Extract area from cloud at given center with distance clearance of shape circle or square"""
if isinstance(full_cloud, np.ndarray):
full_cloud = torch.from_numpy(full_cloud)
if shape == 'square':
x_mask = ((center[0]+clearance) > full_cloud[:, 0]
) & (full_cloud[:, 0] > (center[0]-clearance))
y_mask = ((center[1]+clearance) > full_cloud[:, 1]
) & (full_cloud[:, 1] > (center[1]-clearance))
mask = x_mask & y_mask
elif shape == 'circle':
mask = torch.linalg.norm(full_cloud[:, :2]-center, axis=1) < clearance
else:
raise Exception("Invalid shape")
return mask
def get_voxel(cloud, center, dimensions, return_mask=False):
"""Get voxel or voxel mask from cloud at given center with specified dimensions"""
mask = (cloud[:, :3] >= (center-dimensions/2)
).all(dim=1) & (cloud[:, :3] <= (center+dimensions/2)).all(dim=1)
if return_mask:
return mask
else:
return cloud[mask]
def random_subsample(points, n_samples):
"""Uniform sample of n_samples points from given cloud"""
if points.shape[0] == 0:
print('No points found for random sampling, replacing with dummy')
points = np.zeros((1, points.shape[1]))
# No point sampling if already
if n_samples < points.shape[0]:
random_indices = np.random.choice(
points.shape[0], n_samples, replace=False)
points = points[random_indices, :]
return points
class Early_stop:
def __init__(self, patience=50, min_perc_improvement=0):
self.patience = patience
self.min_perc_improvement = min_perc_improvement
self.not_improved = 0
self.best_loss = torch.tensor(1e+8)
self.last_loss = self.best_loss
self.step = -1
self.loss_hist = []
def log(self, loss):
self.loss_hist.append(loss)
self.step += 1
# Check if improvement by sufficient margin
if (torch.abs(self.last_loss-loss) > torch.abs(self.min_perc_improvement*self.last_loss)) & (self.last_loss < loss):
self.not_improved = 0
else:
self.not_improved += 1
# Keep track of best loss
if loss < self.best_loss:
self.best_loss = loss
# Set last_loss
self.last_loss = loss
if self.not_improved > self.patience:
stop_training = True
else:
stop_training = False
return stop_training
def save_las(pos, path, rgb=None, extra_feature=None, feature_name='Change'):
"""Save a las file with coordinates pos and rgb colors, optional extra features"""
hdr = laspy.header.Header(point_format=2)
outfile = laspy.file.File(path, mode="w", header=hdr)
if not isinstance(extra_feature, type(None)):
outfile.define_new_dimension(
name=feature_name,
data_type=10,
description="Change metric"
)
outfile.writer.set_dimension('change', extra_feature)
allx = pos[:, 0] # Four Points
ally = pos[:, 1]
allz = pos[:, 2]
xmin = np.floor(np.min(allx))
ymin = np.floor(np.min(ally))
zmin = np.floor(np.min(allz))
outfile.header.offset = [xmin, ymin, zmin]
outfile.header.scale = [0.001, 0.001, 0.001]
outfile.x = allx
outfile.y = ally
outfile.z = allz
if not isinstance(rgb, type(None)):
print('Adding color')
if rgb.max() <= 1.0:
rgb *= 255
if rgb.shape[-1] == 3:
outfile.red = rgb[:, 0]
outfile.green = rgb[:, 1]
outfile.blue = rgb[:, 2]
else:
rgb = None
outfile.close()
def co_min_max(tensor_list):
"""Joint min max normalization"""
is_numpy = isinstance(tensor_list[0], np.ndarray)
if is_numpy:
tensor_list = [torch.from_numpy(x) for x in tensor_list]
overall_max = torch.max(torch.stack(
[x[:, :3].max(axis=0)[0] for x in tensor_list]), dim=0)[0]
overall_min = torch.min(torch.stack(
[x[:, :3].min(axis=0)[0] for x in tensor_list]), dim=0)[0]
denominator = overall_max-overall_min + eps
for x in tensor_list:
x[:, :3] = (x[:, :3] - overall_min)/denominator
is_valid(x)
if is_numpy:
tensor_list = [x.numpy() for x in tensor_list]
return tensor_list
def min_max_norm(tensor):
min_ = tensor.min()
max_ = tensor.max()
return (tensor - min_)/(max_ - min_)
def unit_sphere(points, return_inverse=False):
"""Normalize cloud to zero mean and within unit ball"""
mean = points[:, :3].mean(axis=0)
points[:, :3] -= mean
furthest_distance = torch.max(torch.linalg.norm(points[:, :3], dim=-1))
points[:, :3] = points[:, :3] / furthest_distance
if return_inverse:
inverse = {'furthest_distance':furthest_distance,'mean':mean}
return points, inverse
else:
return points
def co_unit_sphere(points_0, points_1, return_inverse=False):
"""Joint zero mean unit ball normalization"""
l_0 = points_0.shape[0]
joint, inverse = unit_sphere(
torch.cat((points_0, points_1)), return_inverse=True)
if return_inverse:
return joint[:l_0, :], joint[l_0:], inverse
else:
return joint[:l_0, :], joint[l_0:]
def view_cloud_o3d(xyz, rgb, show=True):
"""Visualize cloud with o3d"""
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(xyz)
pcd.colors = o3d.utility.Vector3dVector(rgb)
if show:
o3d.visualization.draw_geometries([pcd])
return pcd
def exp_from_paper(x, eps):
"""
Calculate matrix expoential as done in : http://proceedings.mlr.press/v119/xiao20a.html
compute the matrix exponential: \sum_{k=0}^{\infty}\frac{x^{k}}{k!}
"""
scale = int(
np.ceil(np.log2(np.max([torch.norm(x, p=1, dim=-1).max().item(), 0.5]))) + 1)
x = x / (2 ** scale)
s = torch.eye(x.size(-1), device=x.device)
t = x
k = 2
while torch.norm(t, p=1, dim=-1).max().item() > eps:
s = s + t
t = torch.matmul(x, t) / k
k = k + 1
for i in range(scale):
s = torch.matmul(s, s)
return s
def rotate_xy(rad):
"""Create rotation matrix of given rad in 2 dim"""
matrix = torch.tensor(
[[math.cos(rad), -math.sin(rad)], [math.sin(rad), math.cos(rad)]])
def expm(x, eps, algo='torch'):
if algo == 'torch':
return torch.matrix_exp(x)
elif algo == 'original':
return exp_from_paper(x, eps)
else:
raise Exception('Invalid expm algo!')
def rgb_to_hsv(rgb, scale_after=False):
"""Convert rgb values to hsv with optional scaling"""
hsv = torch.zeros_like(rgb)
r = rgb[:, 0]
g = rgb[:, 1]
b = rgb[:, 2]
cmax = rgb.max(axis=1)[0]
cmin = rgb.min(axis=1)[0]
v = cmax
hsv[cmax == cmin, 2] = v[cmax == cmin]
s = (cmax-cmin) / (cmax + eps)
rc = (cmax-r) / (cmax-cmin + eps)
gc = (cmax-g) / (cmax-cmin + eps)
bc = (cmax-b) / (cmax-cmin + eps)
mask_0 = (r == cmax)
mask_1 = (g == cmax)
mask_neither = torch.logical_not(torch.logical_or(mask_0, mask_1))
hsv[mask_0, 0] = (bc-gc)[mask_0]
hsv[mask_1, 0] = (2.0+rc-bc)[mask_1]
hsv[mask_neither, 0] = (4.0+gc-rc)[mask_neither]
hsv[:, 0] = (hsv[:, 0]/6.0) % 1.0
hsv[:, 1] = s
hsv[:, 2] = cmax
if scale_after:
hsv = hsv * torch.Tensor([360, 100, 100])
return hsv
def oversample_cloud(cloud, n_points):
"""Oversample cloud if too few points by repeating the number of missing points to reach the target randomly"""
n_points_original = cloud.shape[0]
if n_points_original >= n_points:
return cloud
else:
random_indices = torch.randint(
0, n_points_original, (n_points - n_points_original,), device=cloud.device)
return torch.cat((cloud, cloud[random_indices, ...]))
def config_loader(path):
"""Load yaml config"""
with open(path) as f:
raw_dict = yaml.load(f,Loader=yaml.FullLoader)
return {key: raw_dict[key]['value'] for key in raw_dict.keys()}
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def sum_except_batch(x, num_dims=1):
'''
Sums all dimensions except the first.
Args:
x: Tensor, shape (batch_size, ...)
num_dims: int, number of batch dims (default=1)
Returns:
x_sum: Tensor, shape (batch_size,)
'''
return x.reshape(*x.shape[:num_dims], -1).sum(-1)
def mean_except_batch(x, num_dims=1):
'''
Averages all dimensions except the first.
Args:
x: Tensor, shape (batch_size, ...)
num_dims: int, number of batch dims (default=1)
Returns:
x_mean: Tensor, shape (batch_size,)
'''
return x.reshape(*x.shape[:num_dims], -1).mean(-1)
def rotate_xy(rad):
matrix = torch.tensor(
[[math.cos(rad), -math.sin(rad)], [math.sin(rad), math.cos(rad)]])
return matrix
def is_valid(tensor):
valid = not torch.logical_or(
tensor.isnan(), tensor.isinf()).any()
return valid
def get_voxel_index(point,min,max,sizes):
axis_size = torch.ceil((max-min) / sizes)
n_total_per_axis = torch.Tensor([axis_size[0:index].prod() for index in range(len(axis_size))])
n_per_axis = ((point - min)//sizes).long()
index = (n_total_per_axis*n_per_axis).sum()
return index
def get_voxel_center(point,min,sizes):
"""Get voxel center given start of voxel,a point from the voxel and the voxel sizes"""
n_per_axis = torch.abs((point - min)//sizes)
min_side = min + (n_per_axis * sizes)
center = min_side + sizes/2
return center
def get_all_voxel_centers(start,end,size):
"""Get all voxel centers given start end (min-max) and voxel sizes"""
n_dims = len(size)
num_voxels = ((end - start) / size).long() + 1
axis_centers = [torch.arange(start[i] + size[i] / 2, end[i] + size[i] / 2, size[i]) for i in range(n_dims)]
centers = torch.stack(torch.meshgrid(*axis_centers[::-1])).reshape((n_dims,-1)).T.flip(-1)
return centers
def voxelize(pos,start,end,size):
"""Get all voxel centers given start end (min-max) and voxel sizes"""
n_dims = len(size)
axis_centers = [torch.arange(start[i] + size[i] / 2, end[i] + size[i] / 2, size[i]) for i in range(n_dims)]
centers = torch.stack(torch.meshgrid(*axis_centers[::-1])).reshape((n_dims,-1)).T.flip(-1)
labels = get_knn(pos,centers,1)
return labels ,centers
def calculate_double_mad(x,c=1.4826):
median = x.median()
abs_dev = (x-median).abs()
left_mad_mask = x<=median
left_median = c*abs_dev[left_mad_mask].median()
right_median = c*abs_dev[~left_mad_mask].median()
x[left_mad_mask] = c * (x - left_median).abs()/left_median
x[~left_mad_mask] = c * (x - right_median).abs()/right_median
return x
def left_mad(x,source_distrib,c=1.4826,cutoff=2.):
median = source_distrib.median()
source_abs_dev = (source_distrib-median).abs()
source_left_mad_mask = source_distrib<=median
left_mad_mask = x<=median
left_mad = c*source_abs_dev[source_left_mad_mask].median()
x[~left_mad_mask] = 0.0
x[left_mad_mask] = (x-median).abs()[left_mad_mask]/left_mad
x[x<cutoff] = 0.0
return x