kcml commited on
Commit
cfce8c8
·
verified ·
1 Parent(s): a02b36e

Upload handcrafted_solution.py

Browse files
Files changed (1) hide show
  1. 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 1 edge.'''
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
- '''Search on sfm depth first. Search on depthmap only if no sfm depth
 
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
- '''If sfm is available in small local region, use it.
207
- Otherwise, search in large region with combined depth'''
208
- #PRINT = False
209
- #if np.isclose(x, 2145.91971831, atol=10) and np.isclose(y, 960.9056338, atol=10):
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 # Cost ->2.6
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 the other vertices'''
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
- ''' entry 0 suggests:
820
- depth_scale = 1
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 # / 2.5 # 2.5 is the scale estimation coefficient # don't use 2.5...
837
- #vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 20.)
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
- print('before del, ', len(vertices), len(vertices_3d), len(connections))
860
- vert_to_del = np.array([-2807.275, -4999.645, 1284.803])
861
- vertices, vertices_3d, connections = delete_one_vert(vertices, vertices_3d, connections, vert_to_del)
862
- print('after del, ', len(vertices), len(vertices_3d), len(connections))
863
- '''
864
  vert_edge_per_image[i] = vertices, connections, vertices_3d
865
- #print(f'detect {len(vertices)} v and {len(connections)} e in view {i}')
866
- #print('before merge, vertices=', vertices, ' connections=', connections, ' vertices_3d=', vertices_3d)
867
- #all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 3.0) # TODO: 3cm looks too small
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: