improved apex detection algorithm

handcrafted_solution.py

@@ -11,11 +11,68 @@ from scipy.spatial.distance import cdist
 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
 
+apex_color = gestalt_color_mapping["apex"]
+eave_end_point = gestalt_color_mapping["eave_end_point"]
+flashing_end_point = gestalt_color_mapping["flashing_end_point"]
+
+apex_color, eave_end_point, flashing_end_point = [np.array(i) for i in [apex_color, eave_end_point, flashing_end_point]]
+unclassified = np.array([(215, 62, 138)])
+line_classes = ['eave', 'ridge', 'rake', 'valley']
 
 def empty_solution():
     '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
     return np.zeros((2,3)), [(0, 1)]
-    
+
+
+
+def undesired_objects(image):
+    image = image.astype('uint8')
+    nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(image, connectivity=8)
+    sizes = stats[:, -1]
+    max_label = 1
+    max_size = sizes[1]
+    for i in range(2, nb_components):
+        if sizes[i] > max_size:
+            max_label = i
+            max_size = sizes[i]
+
+    img2 = np.zeros(output.shape)
+    img2[output == max_label] = 1
+    return img2
+
+def clean_image(image_gestalt) -> np.ndarray:
+    # clears image in from of unclassified and disconected components
+    image_gestalt = np.array(image_gestalt)
+    unclassified_mask = cv2.inRange(image_gestalt, unclassified + 0.0, unclassified + 0.8)
+    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)
+
+    image_gestalt[:, :, 0] *= mask
+    image_gestalt[:, :, 1] *= mask
+    image_gestalt[:, :, 2] *= mask
+    return image_gestalt
+
+
+def get_vertices(image_gestalt, *, color_range=4., dialations=3, erosions=1, kernel_size=13):
+
+    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, flashing_end_point + color_range)
+    # eave_end_point_mask = cv2.bitwise_or(eave_end_point_mask, flashing_end_point_mask)
+
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+
+    apex_mask = cv2.morphologyEx(apex_mask, cv2.MORPH_DILATE, kernel, iterations=dialations)
+    apex_mask = cv2.morphologyEx(apex_mask, cv2.MORPH_ERODE, kernel, iterations=erosions)
+
+    eave_end_point_mask = cv2.morphologyEx(eave_end_point_mask, cv2.MORPH_DILATE, kernel, iterations=dialations)
+    eave_end_point_mask = cv2.morphologyEx(eave_end_point_mask, cv2.MORPH_ERODE, kernel, iterations=erosions)
+
+    *_, apex_centroids = cv2.connectedComponentsWithStats(apex_mask, connectivity=8)
+    *_, other_centroids = cv2.connectedComponentsWithStats(eave_end_point_mask, connectivity=8)
+
+    return [apex_centroids[1:], other_centroids[1:]]
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -42,34 +99,34 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 50.0):
     vertices = []
     connections = []
     # Apex
-    apex_color = np.array(gestalt_color_mapping['apex'])
-    apex_mask = cv2.inRange(gest_seg_np,  apex_color-0.5, apex_color+0.5)
-    if apex_mask.sum() > 0:
-        output = cv2.connectedComponentsWithStats(apex_mask, 8, cv2.CV_32S)
-        (numLabels, labels, stats, centroids) = output
-        stats, centroids = stats[1:], centroids[1:]
-        
-        for i in range(numLabels-1):
-            vert = {"xy": centroids[i], "type": "apex"}
-            vertices.append(vert)
-    
-    eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
-    eave_end_mask = cv2.inRange(gest_seg_np,  eave_end_color-0.5, eave_end_color+0.5)
-    if eave_end_mask.sum() > 0:
-        output = cv2.connectedComponentsWithStats(eave_end_mask, 8, cv2.CV_32S)
-        (numLabels, labels, stats, centroids) = output
-        stats, centroids = stats[1:], centroids[1:]
-        
-        for i in range(numLabels-1):
-            vert = {"xy": centroids[i], "type": "eave_end_point"}
-            vertices.append(vert) 
+    gest_seg_np = clean_image(gest_seg_np)
+    apex_centroids, eave_end_point_centroids = get_vertices(gest_seg_np)
+    # apex_color = np.array(gestalt_color_mapping['apex'])
+    # apex_mask = cv2.inRange(gest_seg_np,  apex_color-0.5, apex_color+0.5)
+    # if apex_mask.sum() > 0:
+    #     output = cv2.connectedComponentsWithStats(apex_mask, 8, cv2.CV_32S)
+    #     (numLabels, labels, stats, centroids) = output
+    #     stats, centroids = stats[1:], centroids[1:]
+    #
+    #     for i in range(numLabels-1):
+    #         vert = {"xy": centroids[i], "type": "apex"}
+    #         vertices.append(vert)
+    #
+    # eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
+    # eave_end_mask = cv2.inRange(gest_seg_np,  eave_end_color-0.5, eave_end_color+0.5)
+    # if eave_end_mask.sum() > 0:
+    #     output = cv2.connectedComponentsWithStats(eave_end_mask, 8, cv2.CV_32S)
+    #     (numLabels, labels, stats, centroids) = output
+    #     stats, centroids = stats[1:], centroids[1:]
+    #
+    #     for i in range(numLabels-1):
+    #         vert = {"xy": centroids[i], "type": "eave_end_point"}
+    #         vertices.append(vert)
     # Connectivity
-    apex_pts = []
-    apex_pts_idxs = []
-    for j, v in enumerate(vertices):
-        apex_pts.append(v['xy'])
-        apex_pts_idxs.append(j)
-    apex_pts = np.array(apex_pts)
+    # apex_pts = []
+    # for j, v in enumerate(vertices):
+    #     apex_pts.append(v['xy'])
+    apex_pts = np.concatenate([apex_centroids, eave_end_point_centroids])
             
     # Ridge connects two apex points
     for edge_class in ['eave', 'ridge', 'rake', 'valley']:
@@ -106,6 +163,8 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 50.0):
                     for a_i, a in enumerate(connected_verts):
                         for b in connected_verts[a_i+1:]:
                             connections.append((a, b))
+    vertices = [{"xy": v, "type": "apex"} for v in apex_centroids]
+    vertices += [{"xy": v, "type": "eave_end_point"} for v in eave_end_point_centroids]
     return vertices, connections
 
 def get_uv_depth(vertices, depth):