algorithm update: in order to be fit, the detected line is required to be in the same direction as proposed one

handcrafted_solution.py

@@ -104,45 +104,38 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=50.0):
     color_range = 4.
     connections = []
     edge_th = edge_th ** 2
+    deviation_threshold = np.cos(np.deg2rad(5))
 
     apex_centroids, eave_end_point_centroids, apex_mask, eave_end_point_mask = get_vertices(gest_seg_np)
 
-    apex_pts = np.concatenate([apex_centroids, eave_end_point_centroids])
+    vertices = np.concatenate([apex_centroids, eave_end_point_centroids])
 
-    # vertex_mask = np.zeros_like(apex_mask)
-    # for i in apex_centroids:
-    #     cv2.circle(vertex_mask, np.round(i).astype(np.uint32), 20, (255,), 40, -1)
-    # for i in eave_end_point_centroids:
-    #     cv2.circle(vertex_mask, np.round(i).astype(np.uint32), 15, (255,), 30, -1)
-    # vertex_mask = np.bitwise_not(vertex_mask)
 
     scale = 1
-    vertex_size = np.zeros(apex_pts.shape[0])
-    for i, coords in enumerate(apex_pts):
+    vertex_size = np.zeros(vertices.shape[0])
+    for i, coords in enumerate(vertices):
         # coords = np.round(coords).astype(np.uint32)
         radius = 25#np.clip(int(max_depth//2 + depth_np[coords[1], coords[0]]), 10, 30)#int(np.clip(max_depth - depth_np[coords[1], coords[0]], 10, 20))
         vertex_size[i] = (scale*radius) ** 2 # because we are using squared distances
-        # cv2.circle(vertex_mask, coords, radius - 5, (255, ), 2 * (radius - 5), -1)
-    # vertex_size *= scale
+
     for edge_class in ['eave', 'ridge', 'rake', 'valley', 'flashing', 'step_flashing']:
-        if len(apex_pts) < 2:
+        if len(vertices) < 2:
             break
         edge_color = np.array(gestalt_color_mapping[edge_class])
 
         mask = cv2.inRange(gest_seg_np,
                            edge_color-color_range,
                            edge_color+color_range)
-        # mask = cv2.bitwise_and(mask, vertex_mask)
-
         mask = cv2.morphologyEx(mask,
                                 cv2.MORPH_DILATE, np.ones((3, 3)), iterations=1)
-        # if edge_class == "ridge":
-        #     mask = cv2.morphologyEx(mask,
-        #                             cv2.MORPH_DILATE, np.ones((11, 11)), iterations=1)
+
 
 
         if np.any(mask):
 
+
+
+
             rho = 1  # distance resolution in pixels of the Hough grid
             theta = np.pi / 180  # angular resolution in radians of the Hough grid
             threshold = 20  # minimum number of votes (intersections in Hough grid cell)
@@ -155,59 +148,108 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=50.0):
             lines = cv2.HoughLinesP(mask, rho, theta, threshold, np.array([]),
                                     min_line_length, max_line_gap)
 
+
             edges = []
 
-            for line in lines if lines is not None else []:
+            if lines is None:
+                continue
+
+            line_directions = np.zeros((len(lines), 2))
+            for line_idx, line in enumerate(lines):
                 for x1,y1,x2,y2 in line:
-                    extend = 40
+                    extend = 35
                     if x1 < x2:
                         x1, y1, x2, y2 = x2, y2, x1, y1
-                    # cv2.line(mask_copy,(x1,y1),(x2,y2),(255,),10, cv2.LINE_AA)
                     direction = (np.array([x2 - x1, y2 - y1]))
-                    direction = extend * direction / np.linalg.norm(direction)
+                    direction = direction / np.linalg.norm(direction)
+                    line_directions[line_idx] = direction
+
+                    direction = extend * direction
 
-                    # norm = extend * np.linalg.norm(np.array([y2 - y1, x1 - x2]))
                     x1,y1 = (-direction + (x1, y1)).astype(np.int32)
                     x2,y2 = (+ direction + (x2, y2)).astype(np.int32)
 
-
                     edges.append((x1, y1, x2, y2))
 
-            edges = np.array(edges)
-            if len(edges) < 1:
-                continue
 
-            begin_distances = cdist(apex_pts, edges[:, :2], metric="sqeuclidean")
-            end_distances = cdist(apex_pts, edges[:, 2:], metric="sqeuclidean")
 
-            begin_closest_points = np.argmin(begin_distances, axis=0) # index of the closest point for each edge end point
-            end_closest_points = np.argmin(end_distances, axis=0)
 
-            begin_closest_point_distances = begin_distances[begin_closest_points, np.arange(len(begin_closest_points))]
-            end_closest_point_distances = end_distances[end_closest_points, np.arange(len(end_closest_points))]
+            edges = np.array(edges)
+            if len(edges) <1:
+                continue
+            # calculate the distances between the vertices and the edge ends
+            begin_distances = cdist(vertices, edges[:, :2], metric="sqeuclidean")
+            end_distances = cdist(vertices, edges[:, 2:], metric="sqeuclidean")
 
+            begin_in_range_mask = begin_distances < vertex_size[:, np.newaxis]
+            end_in_range_mask = end_distances < vertex_size[:, np.newaxis]
 
+            in_range_connected_mask = np.logical_and(np.any(begin_in_range_mask, axis=0),
+                                                     np.any(end_in_range_mask, axis=0))
 
-            begin_in_range_mask = begin_closest_point_distances < vertex_size[begin_closest_points]
-            end_in_range_mask = end_closest_point_distances < vertex_size[end_closest_points]
 
             # where both ends are in range
-            in_range_connected_mask = np.logical_and(begin_in_range_mask, end_in_range_mask)
-
-            edge_idxs = np.where(in_range_connected_mask)[0]
-
-            edges = np.array([begin_closest_points[edge_idxs], end_closest_points[edge_idxs]]).T
-            if len(edges) < 1:
-                continue
-            edges = np.sort(edges, axis=1)
-            unique_edges = np.unique(edges, axis=0)
-
-            unique_edges = unique_edges[unique_edges[:, 0] != unique_edges[:, 1]] # remove self loops
-
-            if len(unique_edges) < 1:
+            begin_in_range_mask = np.logical_and(begin_in_range_mask, in_range_connected_mask)
+            end_in_range_mask = np.logical_and(end_in_range_mask, in_range_connected_mask)
+
+            begin_candidates = np.array(np.where(begin_in_range_mask))
+            end_candidates = np.array(np.where(end_in_range_mask))
+
+            # sort the candidates by line index; required for the seamlessnes np.split
+            sorted_begin_indices = np.argsort(begin_candidates[1])
+            sorted_end_indices = np.argsort(end_candidates[1])
+            begin_candidates = begin_candidates[:, sorted_begin_indices]
+            end_candidates = end_candidates[:, sorted_end_indices]
+
+
+            # create all possible connections between begin and end candidates that correspond to a line
+            grouped_begins = np.split(begin_candidates[0], np.unique(begin_candidates[1], return_index=True)[1][1:])
+            grouped_ends = np.split(end_candidates[0], np.unique(end_candidates[1], return_index=True)[1][1:])
+            line_indices = np.unique(begin_candidates[1])
+
+            # create all possible connections between begin and end candidates that correspond to a line
+            begin_vertex_list = []
+            end_vertex_list = []
+            line_idx_list = []
+            for begin_vertex, end_vertex, line_idx in zip(grouped_begins, grouped_ends, line_indices):
+                begin_vertex, end_vertex = np.meshgrid(begin_vertex, end_vertex)
+                begin_vertex_list.extend(begin_vertex.flatten())
+                end_vertex_list.extend(end_vertex.flatten())
+                line_idx_list.extend([line_idx] * len(begin_vertex.flatten()))
+
+            line_idx_list = np.array(line_idx_list)
+            all_connections = np.array([begin_vertex_list, end_vertex_list])
+
+            # decrease the number of possible connections to reduce number of calculations
+            possible_connections = np.unique(all_connections, axis=1)
+            possible_connections = np.sort(possible_connections, axis=0)
+            possible_connections = np.unique(possible_connections, axis=1)
+            possible_connections = possible_connections[:, possible_connections[0, :] != possible_connections[1, :]]
+
+            if possible_connections.shape[1] < 1:
                 continue
-            connections.extend(unique_edges)
 
+            # precalculate the possible direction vectors
+            possible_direction_vectors = vertices[possible_connections[0]] - vertices[possible_connections[1]]
+            possible_direction_vectors = possible_direction_vectors / np.linalg.norm(possible_direction_vectors, axis=1)[:, np.newaxis]
+
+            owned_lines_per_possible_connections = [list() for i in range(possible_connections.shape[1])]
+
+            # assign lines to possible connections
+            for line_idx, i,j in zip(line_idx_list, begin_vertex_list, end_vertex_list):
+                if i == j:
+                    continue
+                i, j = min(i, j), max(i, j)
+                for connection_idx, connection in enumerate(possible_connections.T):
+                    if np.all((i, j) == connection):
+                        owned_lines_per_possible_connections[connection_idx].append(line_idx)
+                        break
+
+            # check if the lines are in the same direction as the possible connection
+            for fitted_line_idx, owned_lines_per_possible_connection in enumerate(owned_lines_per_possible_connections):
+                line_deviations = np.abs(np.dot(line_directions[owned_lines_per_possible_connection], possible_direction_vectors[fitted_line_idx]))
+                if np.any(line_deviations > deviation_threshold):
+                    connections.append(possible_connections[:, fitted_line_idx])
 
     vertices = [{"xy": v, "type": "apex"} for v in apex_centroids]
     vertices += [{"xy": v, "type": "eave_end_point"} for v in eave_end_point_centroids]