improved edge connection. (but not favorable when depth is still bad)

handcrafted_solution.py

@@ -164,7 +164,7 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np):
     vertex_depth = np.array(vertex_depth)
     return uv, vertex_depth
 
-'''
+
 from numba import njit, prange
 @njit(parallel=True)
 def fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W):
@@ -178,7 +178,7 @@ def fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W):
             if DUMP_IMG: 
               sfm_color_np[j, i] = c
   return sfm_depth_np, sfm_color_np
-'''
+''''''
 
 def get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
   '''Project 3D sfm pointcloud to the image plane '''
@@ -207,7 +207,7 @@ def get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
   #checked = 0  
   #print('dim of us uv zs rgb:', len(us), len(vs), len(zs), len(rgb))
   for u,v,z,c in zip(us,vs,zs, rgb):
-    '''
+    
     sfm_depth_np, sfm_color_np = fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W)
     '''
     i_range = range(max(0, u - dilate_r), min(W, u + dilate_r))
@@ -221,7 +221,7 @@ def get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
             sfm_depth_np[j, i] = z
             if DUMP_IMG: 
               sfm_color_np[j, i] = c
-         
+    '''     
               
   #print(f'checked {checked} pts')
 
@@ -239,7 +239,6 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
     '''Get the vertices and edges from the gestalt segmentation mask of the house'''
     vertices = []
     connections = []
-
     color_th =  10.0 # Cost ->2.6
 
     #-------------------------    
@@ -332,10 +331,12 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
         apex_pts.append(v['xy'])
         apex_pts_idxs.append(j)
     apex_pts = np.array(apex_pts)
-            
-    # Ridge connects two apex points    
-    line_img = np.copy(gest_seg_np) * 0
-    for edge_class in ['eave', 'ridge', 'rake', 'valley']:
+    '''        
+    # Ridge connects two apex points
+    def Ridge_connects_two_apex_points(gest_seg_np, color_th, apex_pts, edge_th):
+      conn = []
+      line_img = np.copy(gest_seg_np) * 0
+      for edge_class in ['eave', 'ridge', 'rake', 'valley']:
         edge_color = np.array(gestalt_color_mapping[edge_class])
         mask = cv2.morphologyEx(cv2.inRange(gest_seg_np,
                                             edge_color-color_th,
@@ -368,10 +369,57 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
                     connected_verts = np.where(connectivity_mask[:,edge_idx])[0]
                     for a_i, a in enumerate(connected_verts):
                         for b in connected_verts[a_i+1:]:
-                            connections.append((a, b))
+                            conn.append((a, b))
+      return conn, line_img
+
+    connections, line_img = Ridge_connects_two_apex_points(gest_seg_np, color_th, apex_pts, edge_th)
+    '''
+    def classifyPairs(apex_pts, apex_pts_idxs, gest_seg_np, apex_mask, eave_end_mask):
+      conn = []
+      # Plot all possible connection pixels in one mask
+      mask = cv2.bitwise_or(apex_mask, eave_end_mask)
+      for edge_class in ['eave', 'ridge', 'rake', 'valley', 'step_flashing']:
+        edge_color = np.array(gestalt_color_mapping[edge_class])
+        mask_e = cv2.morphologyEx(cv2.inRange(gest_seg_np,
+                                            edge_color-color_th,
+                                            edge_color+color_th),
+                                cv2.MORPH_DILATE, np.ones((11, 11)))
+        mask = cv2.bitwise_or(mask, mask_e)
+      # try connecting each apir and see if the cost on the mask is too high
+      def count_on_line_segment(mask, x1, y1, x2, y2, num_points=100):
+        #points = []
+        score = 0
+        for t in range(num_points + 1):
+            t /= num_points
+            x = x1 + t * (x2 - x1)
+            y = y1 + t * (y2 - y1)
+            x, y = x.astype(np.int32), y.astype(np.int32)
+            if mask[y,x] >0:
+              score += 1
+            
+        return score/num_points
+            #points.append((x, y))
+        #return points
+      
+      conn_thr = 0.8  # 80% of pixels are connectivity pixels
+      for p1i in apex_pts_idxs:
+        for p2i in apex_pts_idxs:
+          if p1i == p2i:
+            continue
+          score = count_on_line_segment(mask, apex_pts[p1i][0], apex_pts[p1i][1], apex_pts[p2i][0], apex_pts[p2i][1], num_points=100)          
+          #print(f'{p1i}, {p2i}, score = {score}')
+          if score>conn_thr and ((p2i,p1i) not in conn):
+            conn.append((p1i, p2i))
+
+      return conn, mask
+    connections, line_img = classifyPairs(apex_pts, apex_pts_idxs, gest_seg_np, apex_mask, eave_end_mask)
+    
+    #print(f'{len(vertices)} vertices: {vertices}')
+    #print(len(connections), ' connections: ', connections)
     if DUMP_IMG: 
       filename_edges_map = f'edges_map_{rid}.jpg'
-      cv2.imwrite(filename_edges_map, line_img)
+      if 'line_img' in locals():
+        cv2.imwrite(filename_edges_map, line_img)
       
       
     return vertices, connections
@@ -647,15 +695,18 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
         '''     
         
         vert_edge_per_image[i] = vertices, connections, vertices_3d
-    #all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 3.0) # TODO: 3cm looks too small        
+    #all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 3.0) # TODO: 3cm looks too small       
     all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 150)
     #print(f'after merge, {len(all_3d_vertices)} 3d vertices and {len(connections_3d)} 3d connections')
     #all_3d_vertices_clean, connections_3d_clean  = prune_not_connected(all_3d_vertices, connections_3d)    
     all_3d_vertices, connections_3d  = prune_tall(all_3d_vertices, connections_3d, 1000)
+  
     if len(all_3d_vertices)>35:
       all_3d_vertices, connections_3d  = prune_not_connected(all_3d_vertices, connections_3d)    
-    all_3d_vertices_clean, connections_3d_clean = prune_far(all_3d_vertices, connections_3d, prune_dist_thr=prune_dist_thr)
-    
+    if len(all_3d_vertices)>10:
+      all_3d_vertices_clean, connections_3d_clean = prune_far(all_3d_vertices, connections_3d, prune_dist_thr=prune_dist_thr)
+    else:
+      all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d
     #all_3d_vertices_clean, connections_3d_clean  = all_3d_vertices, connections_3d # don't prune -> cost:2.0
     #print(f'after pruning, {len(all_3d_vertices_clean)} 3d clean vertices and {len(connections_3d_clean)} 3d clean connections')    
     if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1: