fixed bad logic in line selection, vectorized the calculations, added hough lines, changed distance to sqeuclidean, rewrote convert_entry_to_human_readable to use match case

handcrafted_solution.py

@@ -78,44 +78,44 @@ def get_vertices(image_gestalt, *, color_range=4., dialations=3, erosions=1, ker
 
 def convert_entry_to_human_readable(entry):
     out = {}
-    already_good = ['__key__', 'wf_vertices', 'wf_edges', 'edge_semantics', 'mesh_vertices', 'mesh_faces',
-                    'face_semantics', 'K', 'R', 't']
+    already_good = {'__key__', 'wf_vertices', 'wf_edges', 'edge_semantics', 'mesh_vertices', 'mesh_faces',
+                    'face_semantics', 'K', 'R', 't'}
     for k, v in entry.items():
         if k in already_good:
             out[k] = v
             continue
-        if k == 'points3d':
-            out[k] = read_points3D_binary(fid=io.BytesIO(v))
-        if k == 'cameras':
-            out[k] = read_cameras_binary(fid=io.BytesIO(v))
-        if k == 'images':
-            out[k] = read_images_binary(fid=io.BytesIO(v))
-        if k in ['ade20k', 'gestalt']:
-            out[k] = [PImage.open(io.BytesIO(x)).convert('RGB') for x in v]
-        if k == 'depthcm':
-            out[k] = [PImage.open(io.BytesIO(x)) for x in entry['depthcm']]
+        match k:
+            case 'points3d':
+                out[k] = read_points3D_binary(fid=io.BytesIO(v))
+            case 'cameras':
+                out[k] = read_cameras_binary(fid=io.BytesIO(v))
+            case 'images':
+                out[k] = read_images_binary(fid=io.BytesIO(v))
+            case 'ade20k' | 'gestalt':
+                out[k] = [PImage.open(io.BytesIO(x)).convert('RGB') for x in v]
+            case 'depthcm':
+                out[k] = [PImage.open(io.BytesIO(x)) for x in entry['depthcm']]
     return out
 
 
 def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=50.0):
     '''Get the vertices and edges from the gestalt segmentation mask of the house'''
     # Apex
-    color_range = 4.
+    color_range = 8.
     connections = []
+    edge_th = edge_th ** 2
 
-    gest_seg_np = clean_image(gest_seg_np)
     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])
+
     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)
+        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)
+        cv2.circle(vertex_mask, np.round(i).astype(np.uint32), 15, (255,), 30, -1)
     vertex_mask = np.bitwise_not(vertex_mask)
 
-    apex_pts = np.concatenate([apex_centroids, eave_end_point_centroids])
-
-
     for edge_class in ['eave', 'ridge', 'rake', 'valley', 'flashing']:
         if (len(apex_pts) < 2):
             break
@@ -124,39 +124,60 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=50.0):
         mask = cv2.inRange(gest_seg_np,
                            edge_color - color_range,
                            edge_color + color_range)
+        mask = cv2.bitwise_and(mask, vertex_mask)
 
-        if np.any(mask):  # does not really make sense to dilate something if it is empty
-            mask = cv2.bitwise_and(mask, vertex_mask)
-
+        mask = cv2.morphologyEx(mask,
+                                cv2.MORPH_DILATE, np.ones((11, 11)), iterations=2)
+        if edge_class == "ridge":
             mask = cv2.morphologyEx(mask,
-                                    cv2.MORPH_DILATE, np.ones((11, 11)), iterations=3)
-            # line_img = np.zeros_like(gest_seg_np)
-            output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
-            (numLabels, labels, stats, centroids) = output
-            # stats, centroids = stats[1:], centroids[1:]
+                                    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 = 15  # minimum number of votes (intersections in Hough grid cell)
+            min_line_length = 40  # minimum number of pixels making up a line
+            max_line_gap = 30  # maximum gap in pixels between connectable line segments
+
+            # Run Hough on edge detected image
+            # Output "lines" is an array containing endpoints of detected line segments
+            lines = cv2.HoughLinesP(mask, rho, theta, threshold, np.array([]),
+                                    min_line_length, max_line_gap)
+
             edges = []
-            for i in range(1, numLabels):
-                y, x = np.where(labels == i)
-                xleft_idx = np.argmin(x)
-                x_left = x[xleft_idx]
-                y_left = y[xleft_idx]
-                xright_idx = np.argmax(x)
-                x_right = x[xright_idx]
-                y_right = y[xright_idx]
-                edges.append((x_left, y_left, x_right, y_right))
-                # cv2.line(line_img, (x_left, y_left), (x_right, y_right), (255, 255, 255), 2)
+
+            for line in lines if lines is not None else []:
+                for x1, y1, x2, y2 in line:
+                    edges.append((x1, y1, x2, y2))
+
             edges = np.array(edges)
             if len(edges) < 1:
                 continue
-            pts_to_edges_dist = np.minimum(cdist(apex_pts, edges[:, :2]), cdist(apex_pts, edges[:, 2:]))
-            connectivity_mask = pts_to_edges_dist <= edge_th
-            edge_connects = connectivity_mask.sum(axis=0)
-            for edge_idx, edgesum in enumerate(edge_connects):
-                if edgesum >= 2:
-                    connected_verts = np.where(connectivity_mask[:, edge_idx])[0]
-                    for a_i, a in enumerate(connected_verts):
-                        for b in connected_verts[a_i + 1:]:
-                            connections.append((a, b))
+
+            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))]
+
+            begin_in_range_mask = begin_closest_point_distances < edge_th
+            end_in_range_mask = end_closest_point_distances < edge_th
+
+            # 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]
+
+            unique_edges = np.unique(np.array([begin_closest_points[edge_idxs], end_closest_points[edge_idxs]]).T,
+                                     axis=0)
+
+            connections.extend(unique_edges)
+
     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