clean gest

handcrafted_solution.py

@@ -251,7 +251,7 @@ def get_smooth_uv_depth(vertices, depth, gest_seg_np, sfm_depth_np, r=5):
     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):
@@ -265,7 +265,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 '''
@@ -294,7 +294,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))
@@ -308,7 +308,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')
@@ -729,14 +729,14 @@ def prune_far(all_3d_vertices, connections_3d, prune_dist_thr=3000):
     
     return all_3d_vertices, connections_3d
 
-def prune_tall(all_3d_vertices, connections_3d, prune_tall_thr=1000):
+def prune_tall_short(all_3d_vertices, connections_3d, prune_tall_thr=1000, prune_short_thr=100):
     '''Prune vertices that has inpractical z'''
     if (len(all_3d_vertices) < 3) or len(connections_3d) < 1:
       return all_3d_vertices, connections_3d
     
     isolated = []
     for i,v in enumerate(all_3d_vertices):
-      if v[2] > prune_tall_thr:
+      if v[2] > prune_tall_thr or v[2] < prune_short_thr: 
         isolated.append(i)
         break    
 
@@ -749,12 +749,26 @@ def prune_tall(all_3d_vertices, connections_3d, prune_tall_thr=1000):
         return all_3d_vertices, connections_3d
       
       for i,v in enumerate(all_3d_vertices):
-        if v[2] > prune_tall_thr:
+        if v[2] > prune_tall_thr or v[2] < prune_short_thr:
           isolated.append(i)
           break
   
     return all_3d_vertices, connections_3d
 
+def clean_gest(gest_seg_np):
+  bg_color = np.array(gestalt_color_mapping['unclassified'])
+  bg_mask = cv2.inRange(gest_seg_np,  bg_color-10, bg_color+10)
+  fg_mask = cv2.bitwise_not(bg_mask)  
+  if fg_mask.sum() > 0:
+    output = cv2.connectedComponentsWithStats(fg_mask, 8, cv2.CV_32S)
+    (numLabels, labels, stats, centroids) = output
+    sizes = stats[1:, -1]  # Get the areas (skip the first entry which is the background)
+    max_area = max(sizes)
+    max_label = np.where(sizes == max_area)[0] + 1  # Add 1 to get the actual label
+    # mask out anything that doesn't belong to the largest component
+    gest_seg_np[labels != max_label] = bg_color
+
+  return gest_seg_np
 
 def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tuple[np.ndarray, List[int]]:
     good_entry = convert_entry_to_human_readable(entry)
@@ -775,13 +789,16 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
         elif i==2: # only visualize view 0,1
           continue 
         
-        if i!=0:
+        if i!=3:
           continue
         '''
         ade_seg = ade.resize(depth.size)
         ade_seg_np = np.array(ade_seg).astype(np.uint8)
         gest_seg = gest.resize(depth.size)
         gest_seg_np = np.array(gest_seg).astype(np.uint8)
+        gest_seg_np = clean_gest(gest_seg_np)
+        #print('gest_seg_np=', gest_seg_np.shape)
+        
         # Metric3D
         depth_np = np.array(depth) / depth_scale # / 2.5 # 2.5 is the scale estimation coefficient # don't use 2.5...
         #vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 20.)
@@ -818,7 +835,8 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
     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)
+    all_3d_vertices, connections_3d  = prune_tall_short(all_3d_vertices, connections_3d, 1000, 0)
+    
   
     if len(all_3d_vertices)>35:
       all_3d_vertices, connections_3d  = prune_not_connected(all_3d_vertices, connections_3d)    
@@ -826,13 +844,13 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
       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
-
+    
     if i%2:    
       connections_3d_clean = [connections_3d_clean[-1], connections_3d_clean[0]]
     else:
       connections_3d_clean = [connections_3d_clean[0]]
     #print('connections_3d_clean=', connections_3d_clean)
-    
+    ''''''
     #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: