Upload handcrafted_solution.py
Browse files- handcrafted_solution.py +37 -191
handcrafted_solution.py
CHANGED
@@ -18,7 +18,7 @@ if DUMP_IMG:
|
|
18 |
from scipy.sparse import random
|
19 |
|
20 |
def empty_solution():
|
21 |
-
'''Return a minimal valid solution, i.e. 2 vertices and
|
22 |
return np.zeros((2,3)), []
|
23 |
|
24 |
|
@@ -43,77 +43,6 @@ def convert_entry_to_human_readable(entry):
|
|
43 |
return out
|
44 |
|
45 |
|
46 |
-
def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 50.0):
|
47 |
-
'''Get the vertices and edges from the gestalt segmentation mask of the house'''
|
48 |
-
vertices = []
|
49 |
-
connections = []
|
50 |
-
# Apex
|
51 |
-
apex_color = np.array(gestalt_color_mapping['apex'])
|
52 |
-
apex_mask = cv2.inRange(gest_seg_np, apex_color-0.5, apex_color+0.5)
|
53 |
-
if apex_mask.sum() > 0:
|
54 |
-
output = cv2.connectedComponentsWithStats(apex_mask, 8, cv2.CV_32S)
|
55 |
-
(numLabels, labels, stats, centroids) = output
|
56 |
-
stats, centroids = stats[1:], centroids[1:]
|
57 |
-
|
58 |
-
for i in range(numLabels-1):
|
59 |
-
vert = {"xy": centroids[i], "type": "apex"}
|
60 |
-
vertices.append(vert)
|
61 |
-
|
62 |
-
eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
|
63 |
-
eave_end_mask = cv2.inRange(gest_seg_np, eave_end_color-0.5, eave_end_color+0.5)
|
64 |
-
if eave_end_mask.sum() > 0:
|
65 |
-
output = cv2.connectedComponentsWithStats(eave_end_mask, 8, cv2.CV_32S)
|
66 |
-
(numLabels, labels, stats, centroids) = output
|
67 |
-
stats, centroids = stats[1:], centroids[1:]
|
68 |
-
|
69 |
-
for i in range(numLabels-1):
|
70 |
-
vert = {"xy": centroids[i], "type": "eave_end_point"}
|
71 |
-
vertices.append(vert)
|
72 |
-
# Connectivity
|
73 |
-
apex_pts = []
|
74 |
-
apex_pts_idxs = []
|
75 |
-
for j, v in enumerate(vertices):
|
76 |
-
apex_pts.append(v['xy'])
|
77 |
-
apex_pts_idxs.append(j)
|
78 |
-
apex_pts = np.array(apex_pts)
|
79 |
-
|
80 |
-
# Ridge connects two apex points
|
81 |
-
for edge_class in ['eave', 'ridge', 'rake', 'valley']:
|
82 |
-
edge_color = np.array(gestalt_color_mapping[edge_class])
|
83 |
-
mask = cv2.morphologyEx(cv2.inRange(gest_seg_np,
|
84 |
-
edge_color-0.5,
|
85 |
-
edge_color+0.5),
|
86 |
-
cv2.MORPH_DILATE, np.ones((11, 11)))
|
87 |
-
line_img = np.copy(gest_seg_np) * 0
|
88 |
-
if mask.sum() > 0:
|
89 |
-
output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
|
90 |
-
(numLabels, labels, stats, centroids) = output
|
91 |
-
stats, centroids = stats[1:], centroids[1:]
|
92 |
-
edges = []
|
93 |
-
for i in range(1, numLabels):
|
94 |
-
y,x = np.where(labels == i)
|
95 |
-
xleft_idx = np.argmin(x)
|
96 |
-
x_left = x[xleft_idx]
|
97 |
-
y_left = y[xleft_idx]
|
98 |
-
xright_idx = np.argmax(x)
|
99 |
-
x_right = x[xright_idx]
|
100 |
-
y_right = y[xright_idx]
|
101 |
-
edges.append((x_left, y_left, x_right, y_right))
|
102 |
-
cv2.line(line_img, (x_left, y_left), (x_right, y_right), (255, 255, 255), 2)
|
103 |
-
edges = np.array(edges)
|
104 |
-
if (len(apex_pts) < 2) or len(edges) <1:
|
105 |
-
continue
|
106 |
-
pts_to_edges_dist = np.minimum(cdist(apex_pts, edges[:,:2]), cdist(apex_pts, edges[:,2:]))
|
107 |
-
connectivity_mask = pts_to_edges_dist <= edge_th
|
108 |
-
edge_connects = connectivity_mask.sum(axis=0)
|
109 |
-
for edge_idx, edgesum in enumerate(edge_connects):
|
110 |
-
if edgesum>=2:
|
111 |
-
connected_verts = np.where(connectivity_mask[:,edge_idx])[0]
|
112 |
-
for a_i, a in enumerate(connected_verts):
|
113 |
-
for b in connected_verts[a_i+1:]:
|
114 |
-
connections.append((a, b))
|
115 |
-
return vertices, connections
|
116 |
-
|
117 |
def get_uv_depth(vertices, depth):
|
118 |
'''Get the depth of the vertices from the depth image'''
|
119 |
uv = []
|
@@ -129,7 +58,6 @@ def get_uv_depth(vertices, depth):
|
|
129 |
|
130 |
def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
|
131 |
'''Get the depth of the vertices from the depth image'''
|
132 |
-
#print(f'max depth = {np.max(depth)}')
|
133 |
uv = []
|
134 |
for v in vertices:
|
135 |
uv.append(v['xy'])
|
@@ -157,9 +85,6 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
|
|
157 |
def get_local_min(x,y, H, W, depth, sfm_depth_np, r=r, PRINT=False):
|
158 |
'''return a smooth version of detph in radius r'''
|
159 |
local_min = 9999999
|
160 |
-
#PRINT = False
|
161 |
-
#if np.isclose(x, 90.13846154, atol=10) and np.isclose(y, 1175.46495726, atol=10):
|
162 |
-
# PRINT = True
|
163 |
i_range = range(max(0, x - r), min(W, x + r))
|
164 |
j_range = range(max(0, y - r), min(H, y + r))
|
165 |
for i in i_range:
|
@@ -167,18 +92,16 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
|
|
167 |
if sfm_depth_np is not None:
|
168 |
if sfm_depth_np[j, i] != 0:
|
169 |
local_min = min(sfm_depth_np[j, i], local_min)
|
170 |
-
#if PRINT: print(sfm_depth_np[j, i])
|
171 |
if PRINT: print(f'({j},{i})sfm:', sfm_depth_np[j, i])
|
172 |
else:
|
173 |
local_min = min(depth[j, i], local_min)
|
174 |
-
#if PRINT: print(f'({j},{i})dm:', depth[j, i])
|
175 |
else:
|
176 |
local_min = min(depth[j, i], local_min)
|
177 |
-
#if PRINT: print(f'({j},{i})dm:', depth[j, i])
|
178 |
return local_min
|
179 |
|
180 |
def get_priotity_local_min(x,y, H, W, depth, sfm_depth_np, r=r):
|
181 |
-
'''
|
|
|
182 |
exists at all in the local region.
|
183 |
'''
|
184 |
PRINT = False
|
@@ -188,11 +111,8 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
|
|
188 |
xslice = slice(max(0, x - r), min(W, x + r))
|
189 |
local_area = sfm_depth_np[yslice, xslice]
|
190 |
reduced_local_area = local_area[local_area!=0]
|
191 |
-
#if np.isclose(x, 2574.81093605, atol=10) and np.isclose(y, 1063.7265987, atol=10):
|
192 |
-
# PRINT = True
|
193 |
if reduced_local_area.size > 0:
|
194 |
break
|
195 |
-
#print('r=', r)
|
196 |
if reduced_local_area.size > 0:
|
197 |
#print('use sfm')
|
198 |
if PRINT: print(reduced_local_area)
|
@@ -203,23 +123,18 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
|
|
203 |
return get_local_min(x,y, H, W, depth, sfm_depth_np, r, PRINT)
|
204 |
|
205 |
def get_local_min_progressive(x,y, H, W, depth, sfm_depth_np, r=r):
|
206 |
-
'''
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
# PRINT = True
|
211 |
small_r, large_r = 5, 75
|
212 |
PRINT= False
|
213 |
-
#if np.isclose(x, 1799.44303797, atol=10) and np.isclose(y, 622.37721519, atol=10):
|
214 |
-
# PRINT= True
|
215 |
r = small_r
|
216 |
yslice = slice(max(0, y - r), min(H, y + r))
|
217 |
xslice = slice(max(0, x - r), min(W, x + r))
|
218 |
if np.any(sfm_depth_np[yslice, xslice] != 0):
|
219 |
-
#print(f'{x}, {y}, has local sfm')
|
220 |
return get_local_min(x,y, H, W, depth, sfm_depth_np, r)
|
221 |
else:
|
222 |
-
#print(f'{x}, {y}, has local sfm')
|
223 |
r = large_r
|
224 |
local_min = 9999999
|
225 |
i_range = range(max(0, x - r), min(W, x + r))
|
@@ -252,7 +167,7 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
|
|
252 |
vertex_depth = np.array(vertex_depth)
|
253 |
return uv, vertex_depth
|
254 |
|
255 |
-
'''
|
256 |
from numba import njit, prange
|
257 |
@njit(parallel=True)
|
258 |
def fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W):
|
@@ -273,29 +188,20 @@ def get_SfM_depth(XYZ, rgb, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
|
|
273 |
H, W = depth_np.shape[:2]
|
274 |
sfm_depth_np = np.zeros(depth_np.shape)
|
275 |
sfm_color_np = np.zeros(gest_seg_np.shape)
|
276 |
-
#XYZ1 = np.stack([(p.xyz, 1) for p in points3D.values()])
|
277 |
-
#XYZ = np.stack([p.xyz for p in points3D.values()])
|
278 |
-
#rgb = np.stack([p.rgb for p in points3D.values()])
|
279 |
-
#print('XYZ is ', XYZ.shape)
|
280 |
XYZ1 = np.concatenate((XYZ, np.ones((len(XYZ), 1))), axis=1)
|
281 |
-
#print('XYZ1 is ', XYZ1.shape)
|
282 |
Rt = np.concatenate( (R, t.reshape((3,1))), axis=1)
|
283 |
world_to_cam = K @ Rt
|
284 |
xyz = world_to_cam @ XYZ1.transpose()
|
285 |
xyz = np.transpose(xyz)
|
286 |
-
#print('shape of xyz: ', xyz.shape)
|
287 |
valid_idx = ~np.isclose(xyz[:,2], 0, atol=1e-2) & ~np.isnan(xyz[:,0]) & ~np.isnan(xyz[:,1]) & ~np.isnan(xyz[:,2])
|
288 |
xyz = xyz[valid_idx, :]
|
289 |
us, vs, zs = xyz[:,0]/xyz[:,2], xyz[:,1]/xyz[:,2], xyz[:,2]
|
290 |
-
#print('dim of us uv zs rgb:', len(us), len(vs), len(zs), len(rgb))
|
291 |
us = us[~np.isnan(us)]
|
292 |
vs = vs[~np.isnan(vs)]
|
293 |
us = us.astype(np.int32)
|
294 |
vs = vs.astype(np.int32)
|
295 |
-
#checked = 0
|
296 |
-
#print('dim of us uv zs rgb:', len(us), len(vs), len(zs), len(rgb))
|
297 |
for u,v,z,c in zip(us,vs,zs, rgb):
|
298 |
-
'''
|
299 |
sfm_depth_np, sfm_color_np = fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W)
|
300 |
'''
|
301 |
i_range = range(max(0, u - dilate_r), min(W, u + dilate_r))
|
@@ -309,11 +215,7 @@ def get_SfM_depth(XYZ, rgb, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
|
|
309 |
sfm_depth_np[j, i] = z
|
310 |
if DUMP_IMG:
|
311 |
sfm_color_np[j, i] = c
|
312 |
-
|
313 |
-
|
314 |
-
|
315 |
-
#print(f'checked {checked} pts')
|
316 |
-
|
317 |
if DUMP_IMG:
|
318 |
filename_sfm_depth = 'sfm_depth.png'
|
319 |
cv2.imwrite(filename_sfm_depth, sfm_depth_np/100)
|
@@ -328,12 +230,11 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
|
|
328 |
'''Get the vertices and edges from the gestalt segmentation mask of the house'''
|
329 |
vertices = []
|
330 |
connections = []
|
331 |
-
color_th = 10.0
|
332 |
|
333 |
#-------------------------
|
334 |
# combined map from ade
|
335 |
if DUMP_IMG:
|
336 |
-
#print(gest_seg_np.shape, ade_seg_np.shape)
|
337 |
ade_color0 = np.array([0,0,0])
|
338 |
ade_mask0 = cv2.inRange(ade_seg_np, ade_color0-0.5, ade_color0+0.5)
|
339 |
ade_color1 = np.array([120,120,120])
|
@@ -344,26 +245,21 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
|
|
344 |
ade_mask3 = cv2.inRange(ade_seg_np, ade_color3-0.5, ade_color3+0.5)
|
345 |
ade_mask = cv2.bitwise_or(ade_mask3, ade_mask2)
|
346 |
ade_mask = cv2.bitwise_or(ade_mask1, ade_mask)
|
347 |
-
#print(ade_mask.any())
|
348 |
apex_map = np.zeros(ade_seg_np.shape)
|
349 |
apex_map_on_ade = ade_seg_np
|
350 |
apex_map_on_gest = gest_seg_np
|
351 |
# Apex
|
352 |
apex_color = np.array(gestalt_color_mapping['apex'])
|
353 |
-
#print(f'apex_color= {apex_color}')
|
354 |
-
#apex_mask = cv2.inRange(gest_seg_np, apex_color-0.5, apex_color+0.5)
|
355 |
apex_mask = cv2.inRange(gest_seg_np, apex_color-color_th, apex_color+color_th) # include more pts
|
356 |
#apex_mask = cv2.bitwise_and(apex_mask, ade_mask) # remove pts
|
357 |
if apex_mask.sum() > 0:
|
358 |
output = cv2.connectedComponentsWithStats(apex_mask, 8, cv2.CV_32S)
|
359 |
(numLabels, labels, stats, centroids) = output
|
360 |
-
stats, centroids = stats[1:], centroids[1:]
|
361 |
-
|
362 |
for i in range(numLabels-1):
|
363 |
vert = {"xy": centroids[i], "type": "apex"}
|
364 |
vertices.append(vert)
|
365 |
if DUMP_IMG:
|
366 |
-
#print(f'centroids[i]={centroids[i]}')
|
367 |
uu = int(centroids[i][1])
|
368 |
vv = int(centroids[i][0])
|
369 |
# plot a cross
|
@@ -376,11 +272,8 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
|
|
376 |
apex_map_on_ade[uu+ss[0], vv+ss[1]] = (255,255,255)
|
377 |
apex_map_on_gest[uu+ss[0], vv+ss[1]] = (255,255,255)
|
378 |
|
379 |
-
|
380 |
eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
|
381 |
-
#eave_end_mask = cv2.inRange(gest_seg_np, eave_end_color-0.5, eave_end_color+0.5)
|
382 |
eave_end_mask = cv2.inRange(gest_seg_np, eave_end_color-color_th, eave_end_color+color_th)
|
383 |
-
#eave_end_mask = cv2.bitwise_and(eave_end_mask, ade_mask)
|
384 |
if eave_end_mask.sum() > 0:
|
385 |
output = cv2.connectedComponentsWithStats(eave_end_mask, 8, cv2.CV_32S)
|
386 |
(numLabels, labels, stats, centroids) = output
|
@@ -435,7 +328,6 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
|
|
435 |
filename_apex_map = f'apex_map_{rid}.jpg'
|
436 |
cv2.imwrite(filename_apex_map, apex_map)
|
437 |
|
438 |
-
#print(f'{len(vertices)} vertices detected')
|
439 |
# Connectivity
|
440 |
apex_pts = []
|
441 |
apex_pts_idxs = []
|
@@ -443,6 +335,7 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
|
|
443 |
apex_pts.append(v['xy'])
|
444 |
apex_pts_idxs.append(j)
|
445 |
apex_pts = np.array(apex_pts)
|
|
|
446 |
'''
|
447 |
# Ridge connects two apex points
|
448 |
def Ridge_connects_two_apex_points(gest_seg_np, color_th, apex_pts, edge_th):
|
@@ -556,7 +449,6 @@ def merge_vertices_3d(vert_edge_per_image, th=0.1):
|
|
556 |
connections_3d+=[(x+cur_start,y+cur_start) for (x,y) in connections]
|
557 |
cur_start+=len(vertices_3d)
|
558 |
all_3d_vertices = np.concatenate(all_3d_vertices, axis=0)
|
559 |
-
#print (connections_3d)
|
560 |
distmat = cdist(all_3d_vertices, all_3d_vertices)
|
561 |
types = np.array(types).reshape(-1,1)
|
562 |
same_types = cdist(types, types)
|
@@ -586,16 +478,13 @@ def merge_vertices_3d(vert_edge_per_image, th=0.1):
|
|
586 |
for idx in idxs:
|
587 |
old_idx_to_new[idx] = count
|
588 |
count +=1
|
589 |
-
#print (connections_3d)
|
590 |
new_vertices=np.array(new_vertices)
|
591 |
-
#print (connections_3d)
|
592 |
for conn in connections_3d:
|
593 |
new_con = sorted((old_idx_to_new[conn[0]], old_idx_to_new[conn[1]]))
|
594 |
if new_con[0] == new_con[1]:
|
595 |
continue
|
596 |
if new_con not in new_connections:
|
597 |
new_connections.append(new_con)
|
598 |
-
#print (f'{len(new_vertices)} left after merging {len(all_3d_vertices)} with {th=}')
|
599 |
return new_vertices, new_connections
|
600 |
|
601 |
def prune_not_connected(all_3d_vertices, connections_3d):
|
@@ -636,15 +525,12 @@ def uv_to_v3d(uv, depth_vert, K, R, t):
|
|
636 |
|
637 |
def delete_one_vert(vertices, vertices_3d, connections, vert_to_del):
|
638 |
i = np.where(np.all(abs(vertices_3d - vert_to_del) < 0.01, axis=1))
|
639 |
-
#print(i)
|
640 |
-
#print(len(i[0]))
|
641 |
if len(i[0])==0:
|
642 |
if vertices:
|
643 |
return vertices, vertices_3d, connections
|
644 |
else:
|
645 |
return vertices, vertices_3d, connections
|
646 |
|
647 |
-
#print('to del idx=', i[0])
|
648 |
idx = i[0]#[0]
|
649 |
if vertices:
|
650 |
vertices = np.delete(vertices, idx)
|
@@ -659,18 +545,15 @@ def delete_one_vert(vertices, vertices_3d, connections, vert_to_del):
|
|
659 |
connections[ic] = (connections[ic][0]-1, connections[ic][1])
|
660 |
if c[1] >= idx:
|
661 |
connections[ic] = (connections[ic][0], connections[ic][1]-1)
|
662 |
-
|
663 |
-
#print(f'del {len(conn_to_del)} connections')
|
664 |
-
|
665 |
connections = connections.tolist()
|
666 |
-
#print(vertices, vertices_3d, connections)
|
667 |
if vertices:
|
668 |
return vertices, vertices_3d, connections
|
669 |
else:
|
670 |
return vertices_3d, connections
|
671 |
|
672 |
def prune_far(all_3d_vertices, connections_3d, prune_dist_thr=3000):
|
673 |
-
'''Prune vertices that are far away from any
|
674 |
if (len(all_3d_vertices) < 3) or len(connections_3d) < 1:
|
675 |
return all_3d_vertices, connections_3d
|
676 |
|
@@ -678,11 +561,8 @@ def prune_far(all_3d_vertices, connections_3d, prune_dist_thr=3000):
|
|
678 |
distmat = cdist(all_3d_vertices, all_3d_vertices)
|
679 |
for i, v in enumerate(distmat):
|
680 |
exclude_self = np.array([x for idx,x in enumerate(v) if idx!=i])
|
681 |
-
#print('excluded:', exclude_self)
|
682 |
-
#if np.any(exclude_self > prune_dist_thr):
|
683 |
exclude_self = abs(exclude_self)
|
684 |
if min(exclude_self) > prune_dist_thr:
|
685 |
-
#print('del a pt w/ dist = ', min(exclude_self))
|
686 |
isolated.append(i)
|
687 |
break
|
688 |
|
@@ -733,6 +613,9 @@ def prune_tall_short(all_3d_vertices, connections_3d, lowest_z, prune_tall_thr=1
|
|
733 |
return all_3d_vertices, connections_3d
|
734 |
|
735 |
def clean_gest(gest_seg_np):
|
|
|
|
|
|
|
736 |
bg_color = np.array(gestalt_color_mapping['unclassified'])
|
737 |
bg_mask = cv2.inRange(gest_seg_np, bg_color-10, bg_color+10)
|
738 |
if bg_mask.sum() == 0 or bg_mask.sum() == gest_seg_np.shape[0]*gest_seg_np.shape[1]:
|
@@ -750,6 +633,9 @@ def clean_gest(gest_seg_np):
|
|
750 |
return gest_seg_np
|
751 |
|
752 |
def clean_PCD(XYZ, rgb):
|
|
|
|
|
|
|
753 |
lowest_z = 0
|
754 |
center_thr = 500
|
755 |
largest_blob_size = 0
|
@@ -761,8 +647,6 @@ def clean_PCD(XYZ, rgb):
|
|
761 |
clust = OPTICS(min_samples=20, max_eps=150, metric='euclidean', cluster_method='dbscan', algorithm='kd_tree').fit(XYZ)
|
762 |
labels = clust.labels_
|
763 |
unique_labels = set(labels)
|
764 |
-
#print('uni label:', len(unique_labels))
|
765 |
-
#core_samples_mask = np.zeros_like(labels, dtype=bool)
|
766 |
retain_class_mask = labels == -2
|
767 |
if len(unique_labels) > 40 or len(unique_labels) == 1:
|
768 |
return XYZ, rgb, lowest_z
|
@@ -774,9 +658,9 @@ def clean_PCD(XYZ, rgb):
|
|
774 |
largest_blob_size = blob_size
|
775 |
largest_blob = k
|
776 |
|
777 |
-
for k in unique_labels:
|
778 |
-
|
779 |
'''
|
|
|
780 |
if k == -1:
|
781 |
retain_class_mask = retain_class_mask | class_member_mask
|
782 |
continue
|
@@ -816,16 +700,12 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
|
|
816 |
good_entry['R'],
|
817 |
good_entry['t']
|
818 |
)):
|
819 |
-
'''
|
820 |
-
|
821 |
-
if i==1:
|
822 |
-
depth_scale = 2.5
|
823 |
-
elif i==2: # only visualize view 0,1
|
824 |
-
continue
|
825 |
-
|
826 |
if i!=3:
|
827 |
continue
|
828 |
'''
|
|
|
829 |
ade_seg = ade.resize(depth.size)
|
830 |
ade_seg_np = np.array(ade_seg).astype(np.uint8)
|
831 |
gest_seg = gest.resize(depth.size)
|
@@ -833,45 +713,23 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
|
|
833 |
gest_seg_np = clean_gest(gest_seg_np)
|
834 |
|
835 |
# Metric3D
|
836 |
-
depth_np = np.array(depth) / depth_scale
|
837 |
-
|
838 |
-
#vertices, connections = get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_th = 20.)
|
839 |
-
vertices, connections = get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_th = 50.)
|
840 |
-
|
841 |
-
'''
|
842 |
-
if (len(vertices) < 2) or (len(connections) < 1):
|
843 |
-
print (f'Not enough vertices ({len(vertices)}) or connections ({len(connections)}) in image {i}')
|
844 |
-
vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
|
845 |
-
continue
|
846 |
-
'''
|
847 |
if (len(vertices) < 1):
|
848 |
vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
|
849 |
continue
|
850 |
-
|
851 |
-
#uv, depth_vert = get_uv_depth(vertices, depth_np)
|
852 |
-
sfm_depth_np = get_SfM_depth(XYZ, rgb, depth_np, gest_seg_np, K, R, t, 5) # Sensitive. 10 is worse than 0 in testset
|
853 |
-
uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np, sfm_depth_np, 75)
|
854 |
-
#uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np, None)
|
855 |
-
#print('uv:', uv, ' d:', depth_vert)
|
856 |
-
vertices_3d = uv_to_v3d(uv, depth_vert, K, R, t)
|
857 |
|
858 |
-
|
859 |
-
|
860 |
-
|
861 |
-
|
862 |
-
print('after del, ', len(vertices), len(vertices_3d), len(connections))
|
863 |
-
'''
|
864 |
vert_edge_per_image[i] = vertices, connections, vertices_3d
|
865 |
-
|
866 |
-
|
867 |
-
#
|
868 |
all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 150)
|
869 |
-
#print(f'after merge, {len(all_3d_vertices)} 3d vertices and {len(connections_3d)} 3d connections')
|
870 |
-
#print(f'after merge, 3d vertices: {all_3d_vertices} and 3d connections: {connections_3d}')
|
871 |
-
#all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d)
|
872 |
#all_3d_vertices, connections_3d = prune_tall_short(all_3d_vertices, connections_3d, lowest_z, 1000, 0)
|
873 |
-
|
874 |
-
'''
|
875 |
if len(all_3d_vertices)>35:
|
876 |
all_3d_vertices, connections_3d = prune_not_connected(all_3d_vertices, connections_3d)
|
877 |
'''
|
@@ -881,19 +739,7 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
|
|
881 |
all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d
|
882 |
|
883 |
connections_3d_clean = []
|
884 |
-
|
885 |
-
if len(connections_3d_clean):
|
886 |
-
if i%2:
|
887 |
-
connections_3d_clean = [connections_3d_clean[-1], connections_3d_clean[0]]
|
888 |
-
else:
|
889 |
-
connections_3d_clean = [connections_3d_clean[0]]
|
890 |
-
else:
|
891 |
-
print(i, ' has no conn!')
|
892 |
-
#print('connections_3d_clean=', connections_3d_clean)
|
893 |
-
'''
|
894 |
-
#all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d # don't prune -> cost:2.0
|
895 |
-
#print(f'after pruning, {len(all_3d_vertices_clean)} 3d clean vertices and {len(connections_3d_clean)} 3d clean connections')
|
896 |
-
if (len(all_3d_vertices_clean) < 2): # or len(connections_3d_clean) < 1:
|
897 |
print (f'Not enough vertices or connections in the 3D vertices')
|
898 |
return (good_entry['__key__'], *empty_solution())
|
899 |
if visualize:
|
|
|
18 |
from scipy.sparse import random
|
19 |
|
20 |
def empty_solution():
|
21 |
+
'''Return a minimal valid solution, i.e. 2 vertices and 0 edge.'''
|
22 |
return np.zeros((2,3)), []
|
23 |
|
24 |
|
|
|
43 |
return out
|
44 |
|
45 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
46 |
def get_uv_depth(vertices, depth):
|
47 |
'''Get the depth of the vertices from the depth image'''
|
48 |
uv = []
|
|
|
58 |
|
59 |
def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
|
60 |
'''Get the depth of the vertices from the depth image'''
|
|
|
61 |
uv = []
|
62 |
for v in vertices:
|
63 |
uv.append(v['xy'])
|
|
|
85 |
def get_local_min(x,y, H, W, depth, sfm_depth_np, r=r, PRINT=False):
|
86 |
'''return a smooth version of detph in radius r'''
|
87 |
local_min = 9999999
|
|
|
|
|
|
|
88 |
i_range = range(max(0, x - r), min(W, x + r))
|
89 |
j_range = range(max(0, y - r), min(H, y + r))
|
90 |
for i in i_range:
|
|
|
92 |
if sfm_depth_np is not None:
|
93 |
if sfm_depth_np[j, i] != 0:
|
94 |
local_min = min(sfm_depth_np[j, i], local_min)
|
|
|
95 |
if PRINT: print(f'({j},{i})sfm:', sfm_depth_np[j, i])
|
96 |
else:
|
97 |
local_min = min(depth[j, i], local_min)
|
|
|
98 |
else:
|
99 |
local_min = min(depth[j, i], local_min)
|
|
|
100 |
return local_min
|
101 |
|
102 |
def get_priotity_local_min(x,y, H, W, depth, sfm_depth_np, r=r):
|
103 |
+
'''
|
104 |
+
Search on sfm depth first. Search on depthmap only if no sfm depth
|
105 |
exists at all in the local region.
|
106 |
'''
|
107 |
PRINT = False
|
|
|
111 |
xslice = slice(max(0, x - r), min(W, x + r))
|
112 |
local_area = sfm_depth_np[yslice, xslice]
|
113 |
reduced_local_area = local_area[local_area!=0]
|
|
|
|
|
114 |
if reduced_local_area.size > 0:
|
115 |
break
|
|
|
116 |
if reduced_local_area.size > 0:
|
117 |
#print('use sfm')
|
118 |
if PRINT: print(reduced_local_area)
|
|
|
123 |
return get_local_min(x,y, H, W, depth, sfm_depth_np, r, PRINT)
|
124 |
|
125 |
def get_local_min_progressive(x,y, H, W, depth, sfm_depth_np, r=r):
|
126 |
+
'''
|
127 |
+
If sfm is available in small local region, use it.
|
128 |
+
Otherwise, search in large region with combined depth
|
129 |
+
'''
|
|
|
130 |
small_r, large_r = 5, 75
|
131 |
PRINT= False
|
|
|
|
|
132 |
r = small_r
|
133 |
yslice = slice(max(0, y - r), min(H, y + r))
|
134 |
xslice = slice(max(0, x - r), min(W, x + r))
|
135 |
if np.any(sfm_depth_np[yslice, xslice] != 0):
|
|
|
136 |
return get_local_min(x,y, H, W, depth, sfm_depth_np, r)
|
137 |
else:
|
|
|
138 |
r = large_r
|
139 |
local_min = 9999999
|
140 |
i_range = range(max(0, x - r), min(W, x + r))
|
|
|
167 |
vertex_depth = np.array(vertex_depth)
|
168 |
return uv, vertex_depth
|
169 |
|
170 |
+
''' Turn on this to speed up if you have numba
|
171 |
from numba import njit, prange
|
172 |
@njit(parallel=True)
|
173 |
def fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W):
|
|
|
188 |
H, W = depth_np.shape[:2]
|
189 |
sfm_depth_np = np.zeros(depth_np.shape)
|
190 |
sfm_color_np = np.zeros(gest_seg_np.shape)
|
|
|
|
|
|
|
|
|
191 |
XYZ1 = np.concatenate((XYZ, np.ones((len(XYZ), 1))), axis=1)
|
|
|
192 |
Rt = np.concatenate( (R, t.reshape((3,1))), axis=1)
|
193 |
world_to_cam = K @ Rt
|
194 |
xyz = world_to_cam @ XYZ1.transpose()
|
195 |
xyz = np.transpose(xyz)
|
|
|
196 |
valid_idx = ~np.isclose(xyz[:,2], 0, atol=1e-2) & ~np.isnan(xyz[:,0]) & ~np.isnan(xyz[:,1]) & ~np.isnan(xyz[:,2])
|
197 |
xyz = xyz[valid_idx, :]
|
198 |
us, vs, zs = xyz[:,0]/xyz[:,2], xyz[:,1]/xyz[:,2], xyz[:,2]
|
|
|
199 |
us = us[~np.isnan(us)]
|
200 |
vs = vs[~np.isnan(vs)]
|
201 |
us = us.astype(np.int32)
|
202 |
vs = vs.astype(np.int32)
|
|
|
|
|
203 |
for u,v,z,c in zip(us,vs,zs, rgb):
|
204 |
+
''' Use this insead if you have numba
|
205 |
sfm_depth_np, sfm_color_np = fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W)
|
206 |
'''
|
207 |
i_range = range(max(0, u - dilate_r), min(W, u + dilate_r))
|
|
|
215 |
sfm_depth_np[j, i] = z
|
216 |
if DUMP_IMG:
|
217 |
sfm_color_np[j, i] = c
|
218 |
+
|
|
|
|
|
|
|
|
|
219 |
if DUMP_IMG:
|
220 |
filename_sfm_depth = 'sfm_depth.png'
|
221 |
cv2.imwrite(filename_sfm_depth, sfm_depth_np/100)
|
|
|
230 |
'''Get the vertices and edges from the gestalt segmentation mask of the house'''
|
231 |
vertices = []
|
232 |
connections = []
|
233 |
+
color_th = 10.0
|
234 |
|
235 |
#-------------------------
|
236 |
# combined map from ade
|
237 |
if DUMP_IMG:
|
|
|
238 |
ade_color0 = np.array([0,0,0])
|
239 |
ade_mask0 = cv2.inRange(ade_seg_np, ade_color0-0.5, ade_color0+0.5)
|
240 |
ade_color1 = np.array([120,120,120])
|
|
|
245 |
ade_mask3 = cv2.inRange(ade_seg_np, ade_color3-0.5, ade_color3+0.5)
|
246 |
ade_mask = cv2.bitwise_or(ade_mask3, ade_mask2)
|
247 |
ade_mask = cv2.bitwise_or(ade_mask1, ade_mask)
|
|
|
248 |
apex_map = np.zeros(ade_seg_np.shape)
|
249 |
apex_map_on_ade = ade_seg_np
|
250 |
apex_map_on_gest = gest_seg_np
|
251 |
# Apex
|
252 |
apex_color = np.array(gestalt_color_mapping['apex'])
|
|
|
|
|
253 |
apex_mask = cv2.inRange(gest_seg_np, apex_color-color_th, apex_color+color_th) # include more pts
|
254 |
#apex_mask = cv2.bitwise_and(apex_mask, ade_mask) # remove pts
|
255 |
if apex_mask.sum() > 0:
|
256 |
output = cv2.connectedComponentsWithStats(apex_mask, 8, cv2.CV_32S)
|
257 |
(numLabels, labels, stats, centroids) = output
|
258 |
+
stats, centroids = stats[1:], centroids[1:]
|
|
|
259 |
for i in range(numLabels-1):
|
260 |
vert = {"xy": centroids[i], "type": "apex"}
|
261 |
vertices.append(vert)
|
262 |
if DUMP_IMG:
|
|
|
263 |
uu = int(centroids[i][1])
|
264 |
vv = int(centroids[i][0])
|
265 |
# plot a cross
|
|
|
272 |
apex_map_on_ade[uu+ss[0], vv+ss[1]] = (255,255,255)
|
273 |
apex_map_on_gest[uu+ss[0], vv+ss[1]] = (255,255,255)
|
274 |
|
|
|
275 |
eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
|
|
|
276 |
eave_end_mask = cv2.inRange(gest_seg_np, eave_end_color-color_th, eave_end_color+color_th)
|
|
|
277 |
if eave_end_mask.sum() > 0:
|
278 |
output = cv2.connectedComponentsWithStats(eave_end_mask, 8, cv2.CV_32S)
|
279 |
(numLabels, labels, stats, centroids) = output
|
|
|
328 |
filename_apex_map = f'apex_map_{rid}.jpg'
|
329 |
cv2.imwrite(filename_apex_map, apex_map)
|
330 |
|
|
|
331 |
# Connectivity
|
332 |
apex_pts = []
|
333 |
apex_pts_idxs = []
|
|
|
335 |
apex_pts.append(v['xy'])
|
336 |
apex_pts_idxs.append(j)
|
337 |
apex_pts = np.array(apex_pts)
|
338 |
+
# Turns out connection is not a priority
|
339 |
'''
|
340 |
# Ridge connects two apex points
|
341 |
def Ridge_connects_two_apex_points(gest_seg_np, color_th, apex_pts, edge_th):
|
|
|
449 |
connections_3d+=[(x+cur_start,y+cur_start) for (x,y) in connections]
|
450 |
cur_start+=len(vertices_3d)
|
451 |
all_3d_vertices = np.concatenate(all_3d_vertices, axis=0)
|
|
|
452 |
distmat = cdist(all_3d_vertices, all_3d_vertices)
|
453 |
types = np.array(types).reshape(-1,1)
|
454 |
same_types = cdist(types, types)
|
|
|
478 |
for idx in idxs:
|
479 |
old_idx_to_new[idx] = count
|
480 |
count +=1
|
|
|
481 |
new_vertices=np.array(new_vertices)
|
|
|
482 |
for conn in connections_3d:
|
483 |
new_con = sorted((old_idx_to_new[conn[0]], old_idx_to_new[conn[1]]))
|
484 |
if new_con[0] == new_con[1]:
|
485 |
continue
|
486 |
if new_con not in new_connections:
|
487 |
new_connections.append(new_con)
|
|
|
488 |
return new_vertices, new_connections
|
489 |
|
490 |
def prune_not_connected(all_3d_vertices, connections_3d):
|
|
|
525 |
|
526 |
def delete_one_vert(vertices, vertices_3d, connections, vert_to_del):
|
527 |
i = np.where(np.all(abs(vertices_3d - vert_to_del) < 0.01, axis=1))
|
|
|
|
|
528 |
if len(i[0])==0:
|
529 |
if vertices:
|
530 |
return vertices, vertices_3d, connections
|
531 |
else:
|
532 |
return vertices, vertices_3d, connections
|
533 |
|
|
|
534 |
idx = i[0]#[0]
|
535 |
if vertices:
|
536 |
vertices = np.delete(vertices, idx)
|
|
|
545 |
connections[ic] = (connections[ic][0]-1, connections[ic][1])
|
546 |
if c[1] >= idx:
|
547 |
connections[ic] = (connections[ic][0], connections[ic][1]-1)
|
548 |
+
|
|
|
|
|
549 |
connections = connections.tolist()
|
|
|
550 |
if vertices:
|
551 |
return vertices, vertices_3d, connections
|
552 |
else:
|
553 |
return vertices_3d, connections
|
554 |
|
555 |
def prune_far(all_3d_vertices, connections_3d, prune_dist_thr=3000):
|
556 |
+
'''Prune vertices that are far away from any other vertices'''
|
557 |
if (len(all_3d_vertices) < 3) or len(connections_3d) < 1:
|
558 |
return all_3d_vertices, connections_3d
|
559 |
|
|
|
561 |
distmat = cdist(all_3d_vertices, all_3d_vertices)
|
562 |
for i, v in enumerate(distmat):
|
563 |
exclude_self = np.array([x for idx,x in enumerate(v) if idx!=i])
|
|
|
|
|
564 |
exclude_self = abs(exclude_self)
|
565 |
if min(exclude_self) > prune_dist_thr:
|
|
|
566 |
isolated.append(i)
|
567 |
break
|
568 |
|
|
|
613 |
return all_3d_vertices, connections_3d
|
614 |
|
615 |
def clean_gest(gest_seg_np):
|
616 |
+
'''
|
617 |
+
Remove all blobs that are not conencted to the largest blob
|
618 |
+
'''
|
619 |
bg_color = np.array(gestalt_color_mapping['unclassified'])
|
620 |
bg_mask = cv2.inRange(gest_seg_np, bg_color-10, bg_color+10)
|
621 |
if bg_mask.sum() == 0 or bg_mask.sum() == gest_seg_np.shape[0]*gest_seg_np.shape[1]:
|
|
|
633 |
return gest_seg_np
|
634 |
|
635 |
def clean_PCD(XYZ, rgb):
|
636 |
+
'''
|
637 |
+
Remove all points that do not belong to the largest cluster
|
638 |
+
'''
|
639 |
lowest_z = 0
|
640 |
center_thr = 500
|
641 |
largest_blob_size = 0
|
|
|
647 |
clust = OPTICS(min_samples=20, max_eps=150, metric='euclidean', cluster_method='dbscan', algorithm='kd_tree').fit(XYZ)
|
648 |
labels = clust.labels_
|
649 |
unique_labels = set(labels)
|
|
|
|
|
650 |
retain_class_mask = labels == -2
|
651 |
if len(unique_labels) > 40 or len(unique_labels) == 1:
|
652 |
return XYZ, rgb, lowest_z
|
|
|
658 |
largest_blob_size = blob_size
|
659 |
largest_blob = k
|
660 |
|
661 |
+
for k in unique_labels:
|
|
|
662 |
'''
|
663 |
+
# -1 is the noise cluster
|
664 |
if k == -1:
|
665 |
retain_class_mask = retain_class_mask | class_member_mask
|
666 |
continue
|
|
|
700 |
good_entry['R'],
|
701 |
good_entry['t']
|
702 |
)):
|
703 |
+
'''
|
704 |
+
debug per view
|
|
|
|
|
|
|
|
|
|
|
705 |
if i!=3:
|
706 |
continue
|
707 |
'''
|
708 |
+
# (1) 2D processing
|
709 |
ade_seg = ade.resize(depth.size)
|
710 |
ade_seg_np = np.array(ade_seg).astype(np.uint8)
|
711 |
gest_seg = gest.resize(depth.size)
|
|
|
713 |
gest_seg_np = clean_gest(gest_seg_np)
|
714 |
|
715 |
# Metric3D
|
716 |
+
depth_np = np.array(depth) / depth_scale
|
717 |
+
vertices, connections = get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_th = 50.)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
718 |
if (len(vertices) < 1):
|
719 |
vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
|
720 |
continue
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
721 |
|
722 |
+
# (2) Use depth
|
723 |
+
sfm_depth_np = get_SfM_depth(XYZ, rgb, depth_np, gest_seg_np, K, R, t, 5)
|
724 |
+
uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np, sfm_depth_np, 75)
|
725 |
+
vertices_3d = uv_to_v3d(uv, depth_vert, K, R, t)
|
|
|
|
|
726 |
vert_edge_per_image[i] = vertices, connections, vertices_3d
|
727 |
+
|
728 |
+
|
729 |
+
# (3) aggregate info collected from all views:
|
730 |
all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 150)
|
|
|
|
|
|
|
731 |
#all_3d_vertices, connections_3d = prune_tall_short(all_3d_vertices, connections_3d, lowest_z, 1000, 0)
|
732 |
+
''' This didn't help the final solution
|
|
|
733 |
if len(all_3d_vertices)>35:
|
734 |
all_3d_vertices, connections_3d = prune_not_connected(all_3d_vertices, connections_3d)
|
735 |
'''
|
|
|
739 |
all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d
|
740 |
|
741 |
connections_3d_clean = []
|
742 |
+
if (len(all_3d_vertices_clean) < 2):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
743 |
print (f'Not enough vertices or connections in the 3D vertices')
|
744 |
return (good_entry['__key__'], *empty_solution())
|
745 |
if visualize:
|