half implemented missing point detection; scale_estimation_coefficient =4

handcrafted_solution.py

@@ -9,6 +9,7 @@ import numpy as np
 from PIL import Image as PImage
 from hoho.color_mappings import gestalt_color_mapping
 from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
+from scipy.spatial import KDTree
 from scipy.spatial.distance import cdist
 
 apex_color = gestalt_color_mapping["apex"]
@@ -48,6 +49,7 @@ def clean_image(image_gestalt) -> np.ndarray:
     unclassified_mask = cv2.bitwise_not(unclassified_mask)
     mask = undesired_objects(unclassified_mask).astype(np.uint8)
     mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((11, 11), np.uint8), iterations=11)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, np.ones((11, 11), np.uint8), iterations=2)
 
     image_gestalt[:, :, 0] *= mask
     image_gestalt[:, :, 1] *= mask
@@ -56,6 +58,7 @@ def clean_image(image_gestalt) -> np.ndarray:
 
 
 def get_vertices(image_gestalt, *, color_range=4., dialations=3, erosions=1, kernel_size=13):
+    ### detects the apex and eave end and flashing end points
     apex_mask = cv2.inRange(image_gestalt, apex_color - color_range, apex_color + color_range)
     eave_end_point_mask = cv2.inRange(image_gestalt, eave_end_point - color_range, eave_end_point + color_range)
     flashing_end_point_mask = cv2.inRange(image_gestalt, flashing_end_point - color_range,
@@ -76,6 +79,35 @@ def get_vertices(image_gestalt, *, color_range=4., dialations=3, erosions=1, ker
     return apex_centroids[1:], other_centroids[1:], apex_mask, eave_end_point_mask
 
 
+def infer_vertices(image_gestalt, *, color_range=4.):
+    ridge_color = np.array(gestalt_color_mapping["ridge"])
+    rake_color = np.array(gestalt_color_mapping["rake"])
+    ridge_mask = cv2.inRange(image_gestalt,
+                             ridge_color - color_range,
+                             ridge_color + color_range)
+    ridge_mask = cv2.morphologyEx(ridge_mask,
+                                  cv2.MORPH_DILATE, np.ones((3, 3)), iterations=4)
+    rake_mask = cv2.inRange(image_gestalt,
+                            rake_color - color_range,
+                            rake_color + color_range)
+    rake_mask = cv2.morphologyEx(rake_mask,
+                                 cv2.MORPH_DILATE, np.ones((3, 3)), iterations=4)
+
+    intersection_mask = cv2.bitwise_and(ridge_mask, rake_mask)
+    intersection_mask = cv2.morphologyEx(intersection_mask, cv2.MORPH_DILATE, np.ones((11, 11)), iterations=3)
+
+    *_, inferred_centroids = cv2.connectedComponentsWithStats(intersection_mask, connectivity=8, stats=cv2.CV_32S)
+
+    return inferred_centroids[1:], intersection_mask
+
+
+def get_missed_vertices(vertices, inferred_centroids, *,  min_missing_distance=200.0, **kwargs):
+    vertices = KDTree(vertices)
+    closest = vertices.query(inferred_centroids, k=1, distance_upper_bound=min_missing_distance)
+    missed_points = inferred_centroids[closest[1] == len(vertices.data)]
+    return missed_points
+
+
 def convert_entry_to_human_readable(entry):
     out = {}
     already_good = {'__key__', 'wf_vertices', 'wf_edges', 'edge_semantics', 'mesh_vertices', 'mesh_faces',
@@ -98,7 +130,8 @@ def convert_entry_to_human_readable(entry):
     return out
 
 
-def get_vertices_and_edges_from_segmentation(gest_seg_np, *, color_range=4., point_radius=30, max_angle=5., extend = 35, **kwargs):
+def get_vertices_and_edges_from_segmentation(gest_seg_np, *, color_range=4., point_radius=30, max_angle=5., extend=35,
+                                             **kwargs):
     '''Get the vertices and edges from the gestalt segmentation mask of the house'''
     # Apex
     connections = []
@@ -107,6 +140,9 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, color_range=4., poi
     apex_centroids, eave_end_point_centroids, apex_mask, eave_end_point_mask = get_vertices(gest_seg_np)
 
     vertices = np.concatenate([apex_centroids, eave_end_point_centroids])
+    # inferred_vertices, inferred_mask = infer_vertices(gest_seg_np)
+    # missed_vertices = get_missed_vertices(vertices, inferred_vertices, **kwargs)
+    # vertices = np.concatenate([vertices, missed_vertices])
 
     scale = 1
     vertex_size = np.zeros(vertices.shape[0])
@@ -239,6 +275,7 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, color_range=4., poi
                     connections.append(possible_connections[:, fitted_line_idx])
 
     vertices = [{"xy": v, "type": "apex"} for v in apex_centroids]
+    # vertices += [{"xy": v, "type": "apex"} for v in missed_vertices]
     vertices += [{"xy": v, "type": "eave_end_point"} for v in eave_end_point_centroids]
     return vertices, connections
 
@@ -257,7 +294,7 @@ def get_uv_depth(vertices, depth):
     return uv, vertex_depth
 
 
-def merge_vertices_3d(vert_edge_per_image, th=0.1):
+def merge_vertices_3d(vert_edge_per_image, merge_th=0.1, **kwargs):
     '''Merge vertices that are close to each other in 3D space and are of same types'''
     all_3d_vertices = []
     connections_3d = []
@@ -274,7 +311,7 @@ def merge_vertices_3d(vert_edge_per_image, th=0.1):
     distmat = cdist(all_3d_vertices, all_3d_vertices)
     types = np.array(types).reshape(-1, 1)
     same_types = cdist(types, types)
-    mask_to_merge = (distmat <= th) & (same_types == 0)
+    mask_to_merge = (distmat <= merge_th) & (same_types == 0)
     new_vertices = []
     new_connections = []
     to_merge = sorted(list(set([tuple(a.nonzero()[0].tolist()) for a in mask_to_merge])))
@@ -335,7 +372,7 @@ def prune_not_connected(all_3d_vertices, connections_3d):
     return np.array(new_verts), connected_out
 
 
-def predict(entry, visualize=False, **kwargs) -> Tuple[np.ndarray, List[int]]:
+def predict(entry, visualize=False, scale_estimation_coefficient=2.5, **kwargs) -> Tuple[np.ndarray, List[int]]:
     good_entry = convert_entry_to_human_readable(entry)
     if 'gestalt' not in good_entry or 'depthcm' not in good_entry or 'K' not in good_entry or 'R' not in good_entry or 't' not in good_entry:
         print('Missing required fields in the entry')
@@ -350,7 +387,7 @@ def predict(entry, visualize=False, **kwargs) -> Tuple[np.ndarray, List[int]]:
         gest_seg = gest.resize(depth.size)
         gest_seg_np = np.array(gest_seg).astype(np.uint8)
         # Metric3D
-        depth_np = np.array(depth) / 2.5  # 2.5 is the scale estimation coefficient
+        depth_np = np.array(depth) / scale_estimation_coefficient  # 2.5 is the scale estimation coefficient
         vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, **kwargs)
         if (len(vertices) < 2) or (len(connections) < 1):
             print(f'Not enough vertices or connections in image {i}')
@@ -370,7 +407,7 @@ def predict(entry, visualize=False, **kwargs) -> Tuple[np.ndarray, List[int]]:
         vertices_3d = cv2.transform(cv2.convertPointsToHomogeneous(vertices_3d_local), cam_to_world)
         vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
         vert_edge_per_image[i] = vertices, connections, vertices_3d
-    all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 3.0)
+    all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, **kwargs)
     all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d)
     if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
         print(f'Not enough vertices or connections in the 3D vertices')

script.py

@@ -127,7 +127,7 @@ if __name__ == "__main__":
     with ProcessPoolExecutor(max_workers=8) as pool:
         results = []
         for i, sample in enumerate(tqdm(dataset)):
-            results.append(pool.submit(predict, sample, visualize=False, point_radius=25, max_angle=15, extend=30))
+            results.append(pool.submit(predict, sample, visualize=False, point_radius=25, max_angle=15, extend=30, merge_th=3.0,  min_missing_distance=1000000.0, scale_estimation_coefficient=4))
         
         for i, result in enumerate(tqdm(results)):
             key, pred_vertices, pred_edges = result.result()