Improves wireframe prediction with edge fitting

Refactors the wireframe prediction pipeline to incorporate edge fitting based on segment proximity.

The updated approach extracts apex and eave-end vertices and fits lines to edge segments based on proximity within a given threshold. This enables the generation of connections between nearby vertices, refining the overall wireframe structure. Additionally, visualization tools were added for debug purposes.

predict.py

@@ -1,6 +1,6 @@
 import numpy as np
 from typing import Tuple, List
-from hoho2025.example_solutions import empty_solution, read_colmap_rec, get_vertices_and_edges_from_segmentation, get_house_mask, fit_scale_robust_median, get_uv_depth, merge_vertices_3d, prune_not_connected, prune_too_far
+from hoho2025.example_solutions import empty_solution, read_colmap_rec, get_vertices_and_edges_from_segmentation, get_house_mask, fit_scale_robust_median, get_uv_depth, merge_vertices_3d, prune_not_connected, prune_too_far, point_to_segment_dist
 from hoho2025.color_mappings import ade20k_color_mapping, gestalt_color_mapping
 from PIL import Image, ImageDraw
 from visu import save_gestalt_with_proj, draw_crosses_on_image
@@ -9,6 +9,7 @@ import pycolmap
 from PIL import Image as PImage
 import cv2
 import open3d as o3d
+from visu import plot_reconstruction_local, plot_wireframe_local, plot_bpo_cameras_from_entry_local
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -412,15 +413,12 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
         gest_seg = gest.resize(depth_size)
         gest_seg_np = np.array(gest_seg).astype(np.uint8)
         
-        pcloud_segmented, pcloud_idxs = extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id, ade_seg, depth, K=K, R=R, t=t)
-        for idx, p3D in enumerate(colmap_rec.points3D.values()):
-            if idx in pcloud_idxs:
-                p3D.color = np.array([255, 0, 0])
-
+        vertices_ours, connections_ours, vertices_3d_ours = our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id, ade_seg, depth, K=K, R=R, t=t, frame=good_entry)
+        #vertices, connections, vertices_3d = vertices_ours, connections_ours, vertices_3d_ours
         # Get 2D vertices and edges first
-        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=20.)
+        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=25.)
 
-        gt_verts = []
+        #gt_verts = []
         #gt_verts, gt_connects, gt_verts3d = get_gt_vertices_and_edges(good_entry, i, depth, colmap_rec, K, R, t, img_id, ade_seg)
         #vertices, connections = gt_verts, gt_connects
 
@@ -446,15 +444,24 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
             continue
             
         # Call the refactored function to get 3D points
-        vertices_3d = create_3d_wireframe_single_image(
-            vertices, connections, depth, colmap_rec, img_id, ade_seg, K, R, t
-        )
+        vertices_3d = create_3d_wireframe_single_image(vertices, connections, depth, colmap_rec, img_id, ade_seg, K, R, t)
         #vertices_3d = gt_verts3d
         # Store original 2D vertices, connections, and computed 3D points
+
+        if False:
+            pcd, geometries = plot_reconstruction_local(None, colmap_rec, points=True, cameras=True, crop_outliers=True)
+            wireframe = plot_wireframe_local(None, good_entry['wf_vertices'], good_entry['wf_edges'], good_entry['wf_classifications'])
+            wireframe2 = plot_wireframe_local(None, vertices_3d_ours, connections_ours, None, color='rgb(255, 0, 0)')
+            wireframe3 = plot_wireframe_local(None, vertices_3d, connections, None, color='rgb(0, 0, 255)')
+            bpo_cams = plot_bpo_cameras_from_entry_local(None, good_entry)
+
+            visu_all = [pcd] + geometries + wireframe + bpo_cams + wireframe2 + wireframe3
+        #o3d.visualization.draw_geometries(visu_all, window_name="3D Reconstruction")
+
         vert_edge_per_image[i] = vertices, connections, vertices_3d
 
     # Visualize colored COLMAP point cloud with Open3D
-
+    '''
     # Create Open3D point cloud from COLMAP reconstruction
     pcd = o3d.geometry.PointCloud()
 
@@ -472,7 +479,7 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
         
         # Visualize the point cloud
         o3d.visualization.draw_geometries([pcd], window_name="COLMAP Point Cloud")
-    
+    '''
     # Merge vertices from all images
     all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 0.5)
     all_3d_vertices_clean, connections_3d_clean  = prune_not_connected(all_3d_vertices, connections_3d, keep_largest=False)
@@ -485,7 +492,7 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
     return all_3d_vertices_clean, connections_3d_clean 
 
 
-def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg, depth, K=None, R=None, t=None):
+def our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id_substring, ade_seg, depth, K=None, R=None, t=None, ):
     """
     Identify apex and eave-end vertices, then detect lines for eave/ridge/rake/valley.
     Also find all COLMAP points that project into apex or eave_end masks.
@@ -504,9 +511,6 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
     eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
     eave_end_mask = cv2.inRange(gest_seg_np, eave_end_color-10, eave_end_color+10)
     
-    # Combined mask for apex and eave_end
-    combined_mask = cv2.bitwise_or(apex_mask, eave_end_mask)
-    
     H, W = gest_seg_np.shape[:2]
 
     # 1) Find the matching COLMAP image to get its associated 3D points
@@ -518,7 +522,7 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
             break
     if found_img is None:
         print(f"Image substring {img_id_substring} not found in COLMAP.")
-        return np.zeros((H, W), dtype=np.float32), False, None
+        return [], [], []
 
     # 2) Gather 3D points that this image sees (according to COLMAP)
     points_xyz_world = []
@@ -529,15 +533,11 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
             points_idxs.append(pid)
     if not points_xyz_world:
         print(f"No 3D points associated with {found_img.name} in COLMAP.")
-        return np.zeros((H, W), dtype=np.float32), False, found_img # Return found_img for consistency
+        return [], [], []
     
     points_xyz_world = np.array(points_xyz_world)  # (N, 3)
     points_idxs = np.array(points_idxs)          # (N,)
-    
-    # 3) Project points_xyz_world to camera coordinates using R, t
-    # points_cam = R @ points_xyz_world.T + t.reshape(3,1)
-    # points_cam = points_cam.T (N,3)
-    # More robustly:
+
     points_xyz_world_h = np.hstack((points_xyz_world, np.ones((points_xyz_world.shape[0], 1)))) # (N, 4)
     
     # World to Camera transformation matrix
@@ -585,7 +585,7 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
 
     if not uv:
         print(f"No points projected into image bounds for {img_id_substring} using K,R,t.")
-        return np.zeros((H, W), dtype=np.float32), False, found_img
+        return [], [], []
 
     house_mask = get_house_mask(ade_seg)
 
@@ -596,6 +596,10 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
     filtered_points_xyz = []
     filtered_point_idxs = []
     filtered_points_color = []
+    filtered_vertices_apex = []
+    filtered_vertices_apex_uv = []
+    filtered_vertices_eave_end = []
+    filtered_vertices_eave_end_uv = []
 
     # Apex
     apex_color = np.array(gestalt_color_mapping['apex'])
@@ -644,6 +648,26 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
                     filtered_point_idxs.extend(points_idxs[final_valid_indices])
                     filtered_points_color.extend([color] * np.sum(depth_filter))
 
+                    # Find the point with lowest depth in the filtered points
+                    if len(final_valid_indices) > 0:
+                        lowest_depth_idx = np.argmin(depths[depth_filter])
+                        lowest_depth_point = final_valid_indices[lowest_depth_idx]
+                        filtered_vertices_apex.append(points_xyz_world[lowest_depth_point])
+                        filtered_points_xyz.append(points_xyz_world[lowest_depth_point])
+                        filtered_point_idxs.append(points_idxs[lowest_depth_point])
+                        filtered_points_color.append(np.array([1., 1., 0.]))
+
+                        # Project the lowest depth point back to image coordinates for visualization
+                        lowest_cam_point = points_cam[lowest_depth_point]
+
+                        # Project to image plane using K
+                        u_proj = (K[0, 0] * lowest_cam_point[0] / lowest_cam_point[2]) + K[0, 2]
+                        v_proj = (K[1, 1] * lowest_cam_point[1] / lowest_cam_point[2]) + K[1, 2]
+
+                        u_proj_int = int(round(u_proj))
+                        v_proj_int = int(round(v_proj))
+
+                        filtered_vertices_apex_uv.append((u_proj_int, v_proj_int))
 
     # Eave end
     eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
@@ -689,6 +713,27 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
                     filtered_point_idxs.extend(points_idxs[final_valid_indices])
                     filtered_points_color.extend([color] * np.sum(depth_filter))
 
+                    # Find the point with lowest depth in the filtered points
+                    if len(final_valid_indices) > 0:
+                        lowest_depth_idx = np.argmin(depths[depth_filter])
+                        lowest_depth_point = final_valid_indices[lowest_depth_idx]
+                        filtered_vertices_eave_end.append(points_xyz_world[lowest_depth_point])
+                        filtered_points_xyz.append(points_xyz_world[lowest_depth_point])
+                        filtered_point_idxs.append(points_idxs[lowest_depth_point])
+                        filtered_points_color.append(np.array([1., 1., 0.]))
+
+                        # Project the lowest depth point back to image coordinates for visualization
+                        lowest_cam_point = points_cam[lowest_depth_point]
+
+                        # Project to image plane using K
+                        u_proj = (K[0, 0] * lowest_cam_point[0] / lowest_cam_point[2]) + K[0, 2]
+                        v_proj = (K[1, 1] * lowest_cam_point[1] / lowest_cam_point[2]) + K[1, 2]
+
+                        u_proj_int = int(round(u_proj))
+                        v_proj_int = int(round(v_proj))
+
+                        filtered_vertices_eave_end_uv.append((u_proj_int, v_proj_int))
+
     '''
     for i, (u, v) in enumerate(uv):
     # Check if this projected point falls within the combined maskvalid_indices
@@ -697,9 +742,103 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
             filtered_points_xyz.append(points_xyz_world[original_idx])
             filtered_point_idxs.append(points_idxs[original_idx])
     '''
-    filtered_points_xyz = np.array(filtered_points_xyz) if filtered_points_xyz else np.empty((0, 3))
-    filtered_point_idxs = np.array(filtered_point_idxs) if filtered_point_idxs else np.empty((0,))
-    filtered_points_color = np.array(filtered_points_color) if filtered_points_color else np.empty((0, 3))
+    filtered_points_xyz = np.array(filtered_points_xyz[::-1]) if filtered_points_xyz else np.empty((0, 3))
+    filtered_point_idxs = np.array(filtered_point_idxs[::-1]) if filtered_point_idxs else np.empty((0,))
+    filtered_points_color = np.array(filtered_points_color[::-1]) if filtered_points_color else np.empty((0, 3))
+    filtered_vertices_apex = np.array(filtered_vertices_apex) if filtered_vertices_apex else np.empty((0, 3))
+    filtered_vertices_eave_end = np.array(filtered_vertices_eave_end) if filtered_vertices_eave_end else np.empty((0, 3))
+
+
+    connections = []
+    edge_classes = ['eave', 'ridge', 'rake', 'valley']
+    edge_th = 25.0  # threshold for proximity to line segments
+    
+    # Combine apex and eave_end vertices and their UV coordinates
+    all_vertices_3d = []
+    all_vertices_uv = []
+    vertex_types = []
+    
+    # Add apex vertices
+    for i, (vertex_3d, vertex_uv) in enumerate(zip(filtered_vertices_apex, filtered_vertices_apex_uv)):
+        all_vertices_3d.append(vertex_3d)
+        all_vertices_uv.append(vertex_uv)
+        vertex_types.append('apex')
+    
+    # Add eave_end vertices
+    for i, (vertex_3d, vertex_uv) in enumerate(zip(filtered_vertices_eave_end, filtered_vertices_eave_end_uv)):
+        all_vertices_3d.append(vertex_3d)
+        all_vertices_uv.append(vertex_uv)
+        vertex_types.append('eave_end')
+    
+    all_vertices_3d = np.array(all_vertices_3d)
+    all_vertices_uv = np.array(all_vertices_uv)
+    
+    if len(all_vertices_uv) < 2:
+        return [], [], []
+    
+    for edge_class in edge_classes:
+        edge_color = np.array(gestalt_color_mapping[edge_class])
+        mask_raw = cv2.inRange(gest_seg_np, edge_color-10, edge_color+10)
+        # Morphological operations to clean up the mask
+        kernel = np.ones((5, 5), np.uint8)
+        mask = cv2.morphologyEx(mask_raw, cv2.MORPH_CLOSE, kernel)
+        if mask.sum() == 0:
+            continue
+
+        # Connected components
+        output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
+        (numLabels, labels, stats, centroids) = output
+        # Skip the background
+        stats, centroids = stats[1:], centroids[1:]
+        label_indices = range(1, numLabels)
+
+        # For each connected component, do a line fit
+        for lbl in label_indices:
+            ys, xs = np.where(labels == lbl)
+            if len(xs) < 2:
+                continue
+            
+            # Fit a line using cv2.fitLine
+            pts_for_fit = np.column_stack([xs, ys]).astype(np.float32)
+            line_params = cv2.fitLine(pts_for_fit, distType=cv2.DIST_L2, 
+                                      param=0, reps=0.01, aeps=0.01)
+            vx, vy, x0, y0 = line_params.ravel()
+            
+            # Find line segment endpoints by projecting points onto the line
+            proj = ((xs - x0)*vx + (ys - y0)*vy)
+            proj_min, proj_max = proj.min(), proj.max()
+            p1 = np.array([x0 + proj_min*vx, y0 + proj_min*vy])
+            p2 = np.array([x0 + proj_max*vx, y0 + proj_max*vy])
+
+            # Find vertices that are close to this line segment
+            if len(all_vertices_uv) < 2:
+                continue
+
+            # Calculate distance from each vertex UV to the line segment
+            dists = []
+            for vertex_uv in all_vertices_uv:
+                dist = point_to_segment_dist(vertex_uv, p1, p2)
+                dists.append(dist)
+            
+            dists = np.array(dists)
+            
+            # Find vertices that are near this line segment
+            near_mask = (dists <= edge_th)
+            near_indices = np.where(near_mask)[0]
+            
+            if len(near_indices) < 2:
+                continue
+
+            # Connect each pair among these near vertices
+            for i in range(len(near_indices)):
+                for j in range(i+1, len(near_indices)):
+                    idx_a = near_indices[i]
+                    idx_b = near_indices[j]
+                    
+                    # Create connection tuple (using sorted indices for consistency)
+                    conn = tuple(sorted((idx_a, idx_b)))
+                    if conn not in connections:
+                        connections.append(conn)
 
     '''
     depth_fitted, depth_sparse, _, col_img = get_fitted_dense_depth(depth, colmap_rec, img_id_substring, ade_seg, K, R, t)
@@ -765,13 +904,20 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
         pcd_depth.points = o3d.utility.Vector3dVector(segmented_points_3d)
         pcd_depth.colors = o3d.utility.Vector3dVector(np.full((len(segmented_points_3d), 3), [0.0, 0.0, 1.0]))
 
-    # Visualize all point clouds
+    # Visualize all point clouds and spheres
     geometries = [pcd_all]
     if len(filtered_points_xyz) > 0:
         geometries.append(pcd_filtered)
     if len(segmented_points_3d) > 0:
         geometries.append(pcd_depth)
 
-    o3d.visualization.draw_geometries(geometries, window_name=f"Combined Point Cloud - {img_id_substring}")
+    #o3d.visualization.draw_geometries(geometries, window_name=f"Combined Point Cloud - {img_id_substring}")
+    # Convert all_vertices_uv and vertex_types to the required format
+    vertices_formatted = []
+    for uv, vertex_type in zip(all_vertices_uv, vertex_types):
+        vertices_formatted.append({
+            'xy': np.array(uv, dtype=float),
+            'type': vertex_type
+        })
 
-    return filtered_points_xyz, filtered_point_idxs
+    return vertices_formatted, connections, all_vertices_3d