code to clean pcd, but roof points seem critical

handcrafted_solution.py

@@ -8,6 +8,8 @@ import cv2
 from typing import Tuple, List
 
 from scipy.spatial.distance import cdist
+from sklearn import metrics
+from sklearn.cluster import DBSCAN
 
 from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
 from hoho.color_mappings import gestalt_color_mapping, ade20k_color_mapping
@@ -267,14 +269,14 @@ def fill_range(u, v, z, dilate_r, c, sfm_depth_np, sfm_color_np, H, W):
   return sfm_depth_np, sfm_color_np
 ''''''
 
-def get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
+def get_SfM_depth(XYZ, rgb, depth_np, gest_seg_np, K, R, t, dilate_r = 5):
   '''Project 3D sfm pointcloud to the image plane '''
   H, W = depth_np.shape[:2]
   sfm_depth_np = np.zeros(depth_np.shape)
   sfm_color_np = np.zeros(gest_seg_np.shape)
   #XYZ1 = np.stack([(p.xyz, 1) for p in points3D.values()])    
-  XYZ = np.stack([p.xyz for p in points3D.values()])
-  rgb = np.stack([p.rgb for p in points3D.values()])
+  #XYZ = np.stack([p.xyz for p in points3D.values()])
+  #rgb = np.stack([p.rgb for p in points3D.values()])
   #print('XYZ is ', XYZ.shape)
   XYZ1 = np.concatenate((XYZ, np.ones((len(XYZ), 1))), axis=1)
   #print('XYZ1 is ', XYZ1.shape)
@@ -770,9 +772,34 @@ def clean_gest(gest_seg_np):
 
   return gest_seg_np
 
+def clean_PCD(XYZ, rgb):
+  db = DBSCAN(eps=150, min_samples=20).fit(XYZ)
+  labels = db.labels_
+  unique_labels = set(labels)
+  core_samples_mask = np.zeros_like(labels, dtype=bool)
+  center_thr = 500
+  retain_class_mask = labels == -2 
+
+  for k in unique_labels:
+    if k == -1:
+      continue
+    class_member_mask = labels == k
+    pt_k = XYZ[class_member_mask]
+    Xmean = np.mean(pt_k[:,0])
+    Ymean = np.mean(pt_k[:,1])
+    if Xmean < center_thr and Ymean < center_thr:
+      retain_class_mask = retain_class_mask | class_member_mask
+  
+  XYZ = XYZ[retain_class_mask]
+  rgb = rgb[retain_class_mask]
+  return XYZ, rgb
+
 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)
     points3D = good_entry['points3d']
+    XYZ = np.stack([p.xyz for p in points3D.values()])
+    rgb = np.stack([p.rgb for p in points3D.values()])
+    #XYZ, rgb = clean_PCD(XYZ, rgb)
     vert_edge_per_image = {}
     for i, (ade, gest, depth, K, R, t) in enumerate(zip(
                                                 good_entry['ade20k'],
@@ -816,7 +843,7 @@ def predict(entry, visualize=False, prune_dist_thr=600, depth_scale=2.5, ) -> Tu
           continue
 
         #uv, depth_vert = get_uv_depth(vertices, depth_np)         
-        sfm_depth_np = get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, 5) # Sensitive. 10 is worse than 0 in testset
+        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
         uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np, sfm_depth_np, 75)
         #uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np, None)        
         #print('uv:', uv, ' d:', depth_vert)