split some code into functions point_radius=40,

max_angle=15,
extend=25,
merge_th=80.0,
min_missing_distance=500.0,
scale_estimation_coefficient=2.54,
clustering_eps=120,
interpolation_radius=10000,
point_radius_scale=0.5,
# dist_coeff=0,
pointcloud_depth_coeff=1.005,

handcrafted_solution.py

@@ -505,49 +505,83 @@ def prune_not_connected(all_3d_vertices, connections_3d):
 
     return np.array(new_verts), connected_out
 
-
-def predict(entry, visualize=False,
-            scale_estimation_coefficient=2.5,
-            clustering_eps=100,
-            dist_coeff=0,
-            pointcloud_depth_coeff = 1,
-            interpolation_radius=200,
-            **kwargs) -> Tuple[np.ndarray, List[int]]:
-    if 'gestalt' not in entry or 'depthcm' not in entry or 'K' not in entry or 'R' not in entry or 't' not in entry:
-        print('Missing required fields in the entry')
-        return (entry['__key__'], *empty_solution())
-    entry = hoho.decode(entry)
-
-    vert_edge_per_image = {}
+def clean_points3d(entry, clustering_eps):
     image_dict = {}
     for k, v in entry["images"].items():
         image_dict[v.name] = v
     points = [v.xyz for k, v in entry["points3d"].items()]
-
+    
     points = np.array(points)
     point_keys = [k for k, v in entry["points3d"].items()]
     point_keys = np.array(point_keys)
-
+    
     # print(len(points))
-
+    
     clustered = DBSCAN(eps=clustering_eps, min_samples=10).fit(points).labels_
     clustered_indices = np.argsort(clustered)
-
+    
     points = points[clustered_indices]
     point_keys = point_keys[clustered_indices]
     clustered = clustered[clustered_indices]
-
+    
     _, cluster_indices = np.unique(clustered, return_index=True)
-
+    
     clustered_points = np.split(points, cluster_indices[1:])
     clustered_keys = np.split(point_keys, cluster_indices[1:])
-
+    
     biggest_cluster_index = np.argmax([len(i) for i in clustered_points])
     biggest_cluster = clustered_points[biggest_cluster_index]
     biggest_cluster_keys = clustered_keys[biggest_cluster_index]
     biggest_cluster_keys = set(biggest_cluster_keys)
-
+    
     points3d_kdtree = KDTree(biggest_cluster)
+    
+    return points3d_kdtree, biggest_cluster_keys, image_dict
+    
+def get_depthmap_from_pointcloud(image, pointcloud, biggest_cluster_keys, R, t):
+    belonging_points3d = []
+    belonging_points2d = []
+    point_indices = np.where(image.point3D_ids != -1)[0]
+    for idx, point_id in zip(point_indices, image.point3D_ids[point_indices]):
+        if point_id in biggest_cluster_keys:
+            belonging_points3d.append(pointcloud[point_id].xyz)
+            belonging_points2d.append(image.xys[idx])
+    
+    if len(belonging_points3d) < 1:
+        print(f'No 3D points in image {image.name}')
+        raise KeyError
+    belonging_points3d = np.array(belonging_points3d)
+    belonging_points2d = np.array(belonging_points2d)
+    # projected2d, _ = cv2.projectPoints(belonging_points3d, R, t, K, dist_coeff)
+    important = np.where(np.all(belonging_points2d >= 0, axis=1))
+    # Normalize the uv to the camera intrinsics
+    world_to_cam = np.eye(4)
+    world_to_cam[:3, :3] = R
+    world_to_cam[:3, 3] = t
+    
+    homo_belonging_points = cv2.convertPointsToHomogeneous(belonging_points3d)
+    depth = cv2.convertPointsFromHomogeneous(cv2.transform(homo_belonging_points, world_to_cam))
+    depth = depth[:, 0, 2]
+    # projected2d = projected2d[:, 0, :]
+    depth = depth[important[0]]
+    # projected2d = projected2d[important[0]]
+    projected2d = belonging_points2d[important[0]]
+    return projected2d, depth
+def predict(entry, visualize=False,
+            scale_estimation_coefficient=2.5,
+            clustering_eps=100,
+            dist_coeff=0,
+            pointcloud_depth_coeff = 1,
+            interpolation_radius=200,
+            **kwargs) -> Tuple[np.ndarray, List[int]]:
+    if 'gestalt' not in entry or 'depthcm' not in entry or 'K' not in entry or 'R' not in entry or 't' not in entry:
+        print('Missing required fields in the entry')
+        return (entry['__key__'], *empty_solution())
+    entry = hoho.decode(entry)
+
+    vert_edge_per_image = {}
+
+    points3d_kdtree, biggest_cluster_keys, image_dict = clean_points3d(entry, clustering_eps)
 
 
     for i, (gest, depthcm, K, R, t, imagekey) in enumerate(zip(entry['gestalt'],
@@ -569,56 +603,25 @@ def predict(entry, visualize=False,
         depth_np = np.array(depthcm) / scale_estimation_coefficient
         uv, depth_vert_from_depth_map = get_uv_depth(vertices, depth_np)
         try:
-            belonging_points3d = []
-            belonging_points2d = []
-            point_indices = np.where(image_dict[imagekey].point3D_ids != -1)[0]
-            for idx, point_id in zip(point_indices, image_dict[imagekey].point3D_ids[point_indices]):
-                if point_id in biggest_cluster_keys:
-                    belonging_points3d.append(entry["points3d"][point_id].xyz)
-                    belonging_points2d.append(image_dict[imagekey].xys[idx])
-
-            if len(belonging_points3d) < 1:
-                print(f'No 3D points in image {i}')
-                vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
-                raise KeyError
-            belonging_points3d = np.array(belonging_points3d)
-            belonging_points2d = np.array(belonging_points2d)
-            # projected2d, _ = cv2.projectPoints(belonging_points3d, R, t, K, dist_coeff)
-            important = np.where(np.all(belonging_points2d >= 0, axis=1))
-            # Normalize the uv to the camera intrinsics
-            world_to_cam = np.eye(4)
-            world_to_cam[:3, :3] = R
-            world_to_cam[:3, 3] = t
-
-            homo_belonging_points = cv2.convertPointsToHomogeneous(belonging_points3d)
-            depth = cv2.convertPointsFromHomogeneous(cv2.transform(homo_belonging_points, world_to_cam))
-            depth = depth[:, 0, 2]
-            # projected2d = projected2d[:, 0, :]
-            depth = depth[important[0]]
-            # projected2d = projected2d[important[0]]
-            projected2d = belonging_points2d[important[0]]
-            # print(projected2d.shape)
-            # print(depth.shape)
-            depth *= pointcloud_depth_coeff
+            image = image_dict[imagekey]
+
+            projected2d, depth = get_depthmap_from_pointcloud(image, entry["points3d"], biggest_cluster_keys, R, t)
             if len(depth) < 1:
                 print(f'No 3D points in image {i}')
-                vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
+                # vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
                 raise KeyError
-            # print(projected2d.shape, depth.shape)
+            depth *= pointcloud_depth_coeff
 
             # interpolator = si.NearestNDInterpolator(projected2d, depth, rescale=True)
             interpolator = NearestNDInterpolatorWithThreshold(projected2d, depth, interpolation_radius)
-            # interpolator = si.LinearNDInterpolator(projected2d, depth, np.nan)
 
             uv = np.array([v['xy'] for v in vertices])
             xi, yi = uv[:, 0], uv[:, 1]
             depth_vert_from_pointcloud = interpolator(xi, yi)
             depthmap_used = False
 
-            # Get the 3D vertices
         except KeyError:
             #Revert to the depthmap
-            # Metric3D
             depthmap_used = True
 
             # Normalize the uv to the camera intrinsics
@@ -630,12 +633,8 @@ def predict(entry, visualize=False,
 
         depth_vert_nan_idxs = None
         if depthmap_used:
-            # norm_factor = np.max(np.linalg.norm(xy_local, axis=1)[..., None])
             depth_vert = depth_vert_from_depth_map
         else:
-            # 1. query detected vertices in projected2d
-            # if the vertex is beyond some radius, use the depthmap
-            # isnt uv
             depth_vert_nan_idxs = np.where(np.isnan(depth_vert_from_pointcloud))[0]
             depth_vert_from_pointcloud[depth_vert_nan_idxs] = depth_vert_from_depth_map[depth_vert_nan_idxs]
             depth_vert = depth_vert_from_pointcloud

script.py

@@ -146,13 +146,13 @@ if __name__ == "__main__":
                                            max_angle=15,
                                            extend=25,
                                            merge_th=80.0,
-                                           min_missing_distance=350.0,
+                                           min_missing_distance=500.0,
                                            scale_estimation_coefficient=2.54,
                                            clustering_eps=120,
                                            interpolation_radius=10000,
                                            point_radius_scale=0.5,
                                            # dist_coeff=0,
-                                           # pointcloud_depth_coeff=1,
+                                           pointcloud_depth_coeff=1.005,
                                            ))
 
             for result in tqdm(results, desc='Results', total=len(results), position=0):