import io import tempfile import zipfile from collections import defaultdict from typing import Tuple, List, Dict import cv2 import numpy as np import pycolmap from PIL import Image as PImage from scipy.spatial.distance import cdist from sklearn.cluster import DBSCAN from hoho2025.color_mappings import ade20k_color_mapping, gestalt_color_mapping def empty_solution(): '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.''' return np.zeros((2,3)), [(0, 1)] def read_colmap_rec(colmap_data): with tempfile.TemporaryDirectory() as tmpdir: with zipfile.ZipFile(io.BytesIO(colmap_data), "r") as zf: zf.extractall(tmpdir) rec = pycolmap.Reconstruction(tmpdir) return rec def _cam_matrix_from_image(img): """Safely extracts R and t from any pycolmap version.""" cfW = img.cam_from_world if callable(cfW): cfW = cfW() try: R = cfW.rotation.matrix() t = cfW.translation except AttributeError: M = np.array(cfW.matrix()) R, t = M[:, :3], M[:, 3] return np.array(R, dtype=np.float64), np.array(t, dtype=np.float64) def convert_entry_to_human_readable(entry): out = {} for k, v in entry.items(): if 'colmap' in k and k != 'pose_only_in_colmap': out['colmap_binary'] = read_colmap_rec(v) elif k in ['wf_vertices', 'wf_edges', 'K', 'R', 't', 'depth']: try: out[k] = np.array(v) except ValueError as e: if "inhomogeneous" in str(e): out[k] = v else: raise e else: out[k] = v out['__key__'] = entry.get('order_id', 'unknown_id') return out def get_house_mask(ade20k_seg): """ Get a mask of the house in the ADE20K segmentation map. """ house_classes_ade20k = [ 'wall', 'house', 'building;edifice', 'door;double;door', 'windowpane;window', ] np_seg = np.array(ade20k_seg) full_mask = np.zeros(np_seg.shape[:2], dtype=np.uint8) for c in house_classes_ade20k: color = np.array(ade20k_color_mapping[c]) mask = cv2.inRange(np_seg, color-0.5, color+0.5) full_mask = np.logical_or(full_mask, mask) return full_mask def point_to_segment_dist(pt, seg_p1, seg_p2): """ Computes the Euclidean distance from pt to the line segment p1->p2. pt, seg_p1, seg_p2: (x, y) as np.ndarray """ if np.allclose(seg_p1, seg_p2): return np.linalg.norm(pt - seg_p1) seg_vec = seg_p2 - seg_p1 pt_vec = pt - seg_p1 seg_len2 = seg_vec.dot(seg_vec) t = max(0, min(1, pt_vec.dot(seg_vec)/seg_len2)) proj = seg_p1 + t*seg_vec return np.linalg.norm(pt - proj) def project_and_filter_colmap_points( colmap_image: pycolmap.Image, colmap_points3D: Dict[int, 'pycolmap.Point3D'], gest_seg_np: np.ndarray, target_seg_colors: Dict[str, np.ndarray], img_height: int, img_width: int, patch_size: int = 25 ) -> Dict[str, List[np.ndarray]]: """ Project COLMAP 3D points to 2D and filter them based on Gestalt segmentation. Returns: Dict mapping class names to lists of 2D points that fall in those segmentation regions. """ projected_points_by_class = {} for point2D in colmap_image.points2D: if point2D.has_point3D(): u, v = point2D.xy[0], point2D.xy[1] if 0 <= u < img_width and 0 <= v < img_height: half_patch = patch_size // 2 v_start = max(0, int(round(v)) - half_patch) v_end = min(img_height, int(round(v)) + half_patch + 1) u_start = max(0, int(round(u)) - half_patch) u_end = min(img_width, int(round(u)) + half_patch + 1) seg_color_patch = gest_seg_np[v_start:v_end, u_start:u_end] for class_name, target_color in target_seg_colors.items(): patch_matches = np.any(np.all(np.abs(seg_color_patch - target_color) <= 1.0, axis=-1)) if patch_matches: if class_name not in projected_points_by_class: projected_points_by_class[class_name] = [] projected_points_by_class[class_name].append(np.array([u, v])) return projected_points_by_class def cluster_projected_points_to_vertices( projected_points: List[np.ndarray], eps: float, min_samples: int ) -> List[np.ndarray]: """ Cluster projected 2D points using DBSCAN to find vertex candidates. Returns: List of cluster centroids as vertex locations. """ if len(projected_points) < min_samples: return [] X = np.array(projected_points) db = DBSCAN(eps=eps, min_samples=min_samples) labels = db.fit_predict(X) vertex_centroids = [] unique_labels = set(labels) if -1 in unique_labels: unique_labels.remove(-1) for label in unique_labels: class_member_mask = (labels == label) cluster_points = X[class_member_mask] centroid = np.mean(cluster_points, axis=0) vertex_centroids.append(centroid) return vertex_centroids def detect_point_class( gest_seg_np: np.ndarray, class_name: str, gestalt_color_mapping: Dict[str, Tuple[int, int, int]] ) -> List[Dict[str, any]]: """ Detect point-like features (vertices) for a given class in the gestalt segmentation. Args: gest_seg_np: Gestalt segmentation image as numpy array class_name: Name of the class to detect (e.g., 'apex', 'eave_end_point') gestalt_color_mapping: Dictionary mapping class names to RGB colors Returns: List of vertex dictionaries with 'xy' coordinates and 'type' """ vertices = [] if class_name not in gestalt_color_mapping: return vertices class_color = np.array(gestalt_color_mapping[class_name]) class_mask = cv2.inRange(gest_seg_np, class_color-0.5, class_color+0.5) if class_mask.sum() > 0: output = cv2.connectedComponentsWithStats(class_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": class_name} vertices.append(vert) return vertices def verify_edge_mask(pt1, pt2, semantic_mask, min_overlap=0.4): """ Draws a line between two points and verifies that it physically lies on top of the neural network's semantic mask. Kills lines that cross empty space. """ canvas = np.zeros_like(semantic_mask) pt1_int = (int(round(pt1[0])), int(round(pt1[1]))) pt2_int = (int(round(pt2[0])), int(round(pt2[1]))) cv2.line(canvas, pt1_int, pt2_int, 255, 3) line_pixels = cv2.countNonZero(canvas) if line_pixels == 0: return False overlap = cv2.bitwise_and(canvas, semantic_mask) overlap_pixels = cv2.countNonZero(overlap) return (overlap_pixels / line_pixels) >= min_overlap def get_vertices_and_edges_from_segmentation( gest_seg_np: np.ndarray, point_class_names: List[str], edge_class_names: List[str], colmap_image: pycolmap.Image = None, colmap_points3D: Dict[int, 'pycolmap.Point3D'] = None, edge_th: float = 25.0, min_3d_points_for_vertex: int = 1, vertex_cluster_eps: float = 5.0, use_colmap_for_vertices: bool = True, patch_size: int = 25 ) -> Tuple[List[dict], List[Tuple[int, int]]]: """ Identify apex and eave-end vertices, then detect lines for eave/ridge/rake/valley. Now enhanced with COLMAP 3D point projection and DBSCAN clustering for vertex detection. """ if not isinstance(gest_seg_np, np.ndarray): gest_seg_np = np.array(gest_seg_np) vertices = [] H, W = gest_seg_np.shape[:2] colmap_width = colmap_image.camera.width if colmap_image is not None else None colmap_height = colmap_image.camera.height if colmap_image is not None else None if colmap_width != W or colmap_height != H: print(f"Warning: colmap image size {colmap_width}x{colmap_height} does not match gestalt segmentation size {W}x{H}") if use_colmap_for_vertices and colmap_image is not None and colmap_points3D is not None: try: target_seg_colors = {} for class_name in point_class_names: if class_name in gestalt_color_mapping: target_seg_colors[class_name] = np.array(gestalt_color_mapping[class_name]) projected_points_by_class = project_and_filter_colmap_points( colmap_image, colmap_points3D, gest_seg_np, target_seg_colors, H, W, patch_size=patch_size ) for class_name in point_class_names: if class_name in projected_points_by_class: points_for_class = projected_points_by_class[class_name] class_centroids = cluster_projected_points_to_vertices( points_for_class, eps=vertex_cluster_eps, min_samples=min_3d_points_for_vertex ) for centroid in class_centroids: vert = {"xy": centroid, "type": class_name} vertices.append(vert) print(f"Found {len(vertices)} vertices using COLMAP projection and clustering") except Exception as e: print(f"Error using COLMAP for vertex detection: {e}") vertices = [] if len(vertices) < 2: print("Using fallback method for vertex detection") vertices = [] for class_name in point_class_names: point_vertices = detect_point_class(gest_seg_np, class_name, gestalt_color_mapping) vertices.extend(point_vertices) structural_pts = [] structural_idx_map = [] for idx, v in enumerate(vertices): structural_pts.append(v['xy']) structural_idx_map.append(idx) structural_pts = np.array(structural_pts) connections = [] for edge_class in edge_class_names: if edge_class in ['ridge']: allowed_types = ['apex'] elif edge_class in ['eave', 'flashing', 'step_flashing']: allowed_types = ['eave_end_point', 'flashing_end_point'] else: allowed_types = ['apex', 'eave_end_point', 'flashing_end_point'] allowed_pts = [] allowed_idx_map = [] for orig_idx, v in enumerate(vertices): if v['type'] in allowed_types: allowed_pts.append(v['xy']) allowed_idx_map.append(orig_idx) allowed_pts = np.array(allowed_pts) if len(allowed_pts) < 2: continue edge_color = np.array(gestalt_color_mapping[edge_class]) mask_raw = cv2.inRange(gest_seg_np, edge_color-0.5, edge_color+0.5) kernel = np.ones((5, 5), np.uint8) mask = cv2.morphologyEx(mask_raw, cv2.MORPH_CLOSE, kernel) if mask.sum() == 0: continue output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S) (numLabels, labels, stats, centroids) = output stats, centroids = stats[1:], centroids[1:] label_indices = range(1, numLabels) for lbl in label_indices: mask_i = np.zeros_like(mask) mask_i[labels == lbl] = 255 lines = cv2.HoughLinesP(mask_i, rho=1, theta=np.pi/180, threshold=15, minLineLength=8, maxLineGap=20) if lines is None: continue for line in lines: x1, y1, x2, y2 = line[0] p1 = np.array([x1, y1], dtype=np.float32) p2 = np.array([x2, y2], dtype=np.float32) if len(allowed_pts) < 2: continue dists = np.array([ point_to_segment_dist(allowed_pts[i], p1, p2) for i in range(len(allowed_pts)) ]) near_mask = (dists <= edge_th) near_indices = np.where(near_mask)[0] if len(near_indices) < 2: continue 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] vA = allowed_idx_map[idx_a] vB = allowed_idx_map[idx_b] conn = tuple(sorted((vA, vB))) if conn not in connections: is_valid_edge = verify_edge_mask(allowed_pts[idx_a], allowed_pts[idx_b], mask, min_overlap=0.3) if is_valid_edge: connections.append(conn) return vertices, connections def get_uv_depth(vertices: List[dict], depth_fitted: np.ndarray, sparse_depth: np.ndarray, search_radius: int = 10) -> Tuple[np.ndarray, np.ndarray]: """For each vertex return its (u, v) and a depth value. Uses the nearest valid sparse-depth pixel within search_radius of the vertex; falls back to the dense depth_fitted value when no sparse depth is available. """ uv = np.array([vert['xy'] for vert in vertices], dtype=np.float32) uv_int = np.round(uv).astype(np.int32) H, W = depth_fitted.shape[:2] uv_int[:, 0] = np.clip(uv_int[:, 0], 0, W - 1) uv_int[:, 1] = np.clip(uv_int[:, 1], 0, H - 1) vertex_depth = np.zeros(len(vertices), dtype=np.float32) dense_count = 0 for i, (x_i, y_i) in enumerate(uv_int): x0 = max(0, x_i - search_radius) x1 = min(W, x_i + search_radius + 1) y0 = max(0, y_i - search_radius) y1 = min(H, y_i + search_radius + 1) region = sparse_depth[y0:y1, x0:x1] valid_y, valid_x = np.where(region > 0) if valid_y.size > 0: global_x = x0 + valid_x global_y = y0 + valid_y dist_sq = (global_x - x_i)**2 + (global_y - y_i)**2 min_idx = np.argmin(dist_sq) vertex_depth[i] = region[valid_y[min_idx], valid_x[min_idx]] else: vertex_depth[i] = depth_fitted[y_i, x_i] dense_count += 1 return uv, vertex_depth def project_vertices_to_3d(uv: np.ndarray, depth_vert: np.ndarray, col_img: pycolmap.Image, colmap_rec: pycolmap.Reconstruction) -> np.ndarray: xy_local = np.ones((len(uv), 3)) try: K = col_img.camera.calibration_matrix() except AttributeError: K = colmap_rec.cameras[col_img.camera_id].calibration_matrix() xy_local[:, 0] = (uv[:, 0] - K[0, 2]) / K[0, 0] xy_local[:, 1] = (uv[:, 1] - K[1, 2]) / K[1, 1] vertices_3d_local = xy_local * depth_vert[...,None] R, t = _cam_matrix_from_image(col_img) world_to_cam = np.eye(4) world_to_cam[:3, :3] = R world_to_cam[:3, 3] = t cam_to_world = np.linalg.inv(world_to_cam) vertices_3d_homogeneous = cv2.convertPointsToHomogeneous(vertices_3d_local) vertices_3d = cv2.transform(vertices_3d_homogeneous, cam_to_world) vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3) return vertices_3d def create_3d_wireframe_single_image(vertices: List[dict], connections: List[Tuple[int, int]], depth: PImage.Image, colmap_rec: pycolmap.Reconstruction, img_id: str, ade_seg: PImage.Image) -> np.ndarray: """Lift one image view's 2D vertices to 3D world coordinates. Fits the dense depth to the sparse COLMAP depth, reads a depth per vertex, and back-projects. Returns an empty (0, 3) array if there is no sparse depth. """ if (len(vertices) < 2) or (len(connections) < 1): print(f'Warning: create_3d_wireframe_single_image called with insufficient vertices/connections for image {img_id}') return np.empty((0, 3)) depth_fitted, depth_sparse, found_sparse, col_img = get_fitted_dense_depth( depth, colmap_rec, img_id, ade_seg ) if not found_sparse or col_img is None: return np.empty((0, 3)) uv, depth_vert = get_uv_depth(vertices, depth_fitted, depth_sparse, search_radius=25) vertices_3d = project_vertices_to_3d(uv, depth_vert, col_img, colmap_rec) return vertices_3d def merge_vertices_3d(vert_edge_per_image, point_class_names: List[str], th=0.5): '''Merge vertices that are close in 3D space and of the same type.''' all_3d_vertices = [] connections_3d = [] all_indexes = [] cur_start = 0 types = [] type_to_id = {class_name: idx for idx, class_name in enumerate(point_class_names)} for cimg_idx, (vertices, connections, vertices_3d) in vert_edge_per_image.items(): vertex_type_ids = [] for v in vertices: vertex_type = v['type'] type_id = type_to_id.get(vertex_type, -1) vertex_type_ids.append(type_id) types += vertex_type_ids all_3d_vertices.append(vertices_3d) connections_3d+=[(x+cur_start,y+cur_start) for (x,y) in connections] cur_start+=len(vertices_3d) all_3d_vertices = np.concatenate(all_3d_vertices, axis=0) distmat = cdist(all_3d_vertices, all_3d_vertices) types = np.array(types).reshape(-1,1) same_types = cdist(types, types) # Merge vertices that are both close in space and of the same type. mask_to_merge = (distmat <= th) & (same_types==0) new_vertices = [] new_connections = [] to_merge = sorted(list(set([tuple(a.nonzero()[0].tolist()) for a in mask_to_merge]))) # Transitive grouping: union overlapping merge-sets into connected groups. to_merge_final = defaultdict(list) for i in range(len(all_3d_vertices)): for j in to_merge: if i in j: to_merge_final[i]+=j for k, v in to_merge_final.items(): to_merge_final[k] = list(set(v)) already_there = set() merged = [] for k, v in to_merge_final.items(): if k in already_there: continue merged.append(v) for vv in v: already_there.add(vv) old_idx_to_new = {} count=0 for idxs in merged: new_vertices.append(all_3d_vertices[idxs].mean(axis=0)) for idx in idxs: old_idx_to_new[idx] = count count +=1 new_vertices=np.array(new_vertices) for conn in connections_3d: new_con = sorted((old_idx_to_new[conn[0]], old_idx_to_new[conn[1]])) if new_con[0] == new_con[1]: continue if new_con not in new_connections: new_connections.append(new_con) return new_vertices, new_connections def prune_not_connected(all_3d_vertices, connections_3d, keep_largest=True): """ Prune vertices not connected to anything. If keep_largest=True, also keep only the largest connected component in the graph. """ if len(all_3d_vertices) == 0: return np.empty((0, 3)), [] adj = defaultdict(set) for (i, j) in connections_3d: adj[i].add(j) adj[j].add(i) used_idxs = set() for (i, j) in connections_3d: used_idxs.add(i) used_idxs.add(j) if not used_idxs: return np.empty((0,3)), [] # If we only want to remove truly isolated points, but keep multiple subgraphs: if not keep_largest: new_map = {} used_list = sorted(list(used_idxs)) for new_id, old_id in enumerate(used_list): new_map[old_id] = new_id new_vertices = np.array([all_3d_vertices[old_id] for old_id in used_list]) new_conns = [] for (i, j) in connections_3d: if i in used_idxs and j in used_idxs: new_conns.append((new_map[i], new_map[j])) return new_vertices, new_conns # Otherwise find the largest connected component: visited = set() def bfs(start): queue = [start] comp = [] visited.add(start) while queue: cur = queue.pop() comp.append(cur) for neigh in adj[cur]: if neigh not in visited: visited.add(neigh) queue.append(neigh) return comp comps = [] for idx in used_idxs: if idx not in visited: c = bfs(idx) comps.append(c) comps.sort(key=lambda c: len(c), reverse=True) largest = comps[0] if len(comps)>0 else [] new_map = {} for new_id, old_id in enumerate(largest): new_map[old_id] = new_id new_vertices = np.array([all_3d_vertices[old_id] for old_id in largest]) new_conns = [] for (i, j) in connections_3d: if i in largest and j in largest: new_conns.append((new_map[i], new_map[j])) new_conns = list(set([tuple(sorted(c)) for c in new_conns])) return new_vertices, new_conns def get_sparse_depth(colmap_rec, img_id_substring, depth): H, W = depth.shape found_img = None for img_id_c, col_img in colmap_rec.images.items(): if img_id_substring in col_img.name: found_img = col_img break if found_img is None: return np.zeros((H, W), dtype=np.float32), False, None points_xyz = [] for pid, p3D in colmap_rec.points3D.items(): if found_img.has_point3D(pid): points_xyz.append(p3D.xyz) if not points_xyz: return np.zeros((H, W), dtype=np.float32), False, found_img points_xyz = np.array(points_xyz) uv = [] z_vals = [] cam = colmap_rec.cameras[found_img.camera_id] R, t = _cam_matrix_from_image(found_img) K = cam.calibration_matrix() for xyz in points_xyz: p_cam = R @ np.asarray(xyz, dtype=np.float64) + t if p_cam[2] > 0: u = p_cam[0] / p_cam[2] * K[0, 0] + K[0, 2] v = p_cam[1] / p_cam[2] * K[1, 1] + K[1, 2] u_i, v_i = int(round(u)), int(round(v)) if 0 <= u_i < W and 0 <= v_i < H: uv.append((u_i, v_i)) z_vals.append(p_cam[2]) uv = np.array(uv, dtype=int) z_vals = np.array(z_vals) depth_out = np.zeros((H, W), dtype=np.float32) if len(uv) > 0: depth_out[uv[:,1], uv[:,0]] = z_vals return depth_out, True, found_img def fit_scale_robust_median(depth, sparse_depth, validity_mask=None): """ Fit a scale factor to the depth map using the median of the ratio of sparse to dense depth. """ if validity_mask is None: mask = (sparse_depth != 0) else: mask = (sparse_depth != 0) & validity_mask mask = mask & (depth <50) & (sparse_depth <50) X = depth[mask] Y = sparse_depth[mask] alpha =np.median(Y/X) depth_fitted = alpha * depth return alpha, depth_fitted def get_fitted_dense_depth(depth, colmap_rec, img_id, ade20k_seg): """Scale the dense depth to align with the sparse COLMAP depth. Reads sparse depth from COLMAP, fits a scale factor using only points inside the ADE20k house mask, and returns the scaled dense depth and the sparse depth. found_sparse is False when no sparse depth is available for the image. """ depth_np = np.array(depth) / 1000. # mm to meters depth_sparse, found_sparse, col_img = get_sparse_depth(colmap_rec, img_id, depth_np) if not found_sparse: print(f'No sparse depth found for image {img_id}') return depth_np, np.zeros_like(depth_np), False, None house_mask = get_house_mask(ade20k_seg) k, depth_fitted = fit_scale_robust_median(depth_np, depth_sparse, validity_mask=house_mask) print(f"Fitted depth scale k={k:.4f} for image {img_id}") return depth_fitted, depth_sparse, True, col_img def precompute_overlapping_views( colmap_reconstruction: pycolmap.Reconstruction, min_shared_points: int = 10 ) -> Dict[int, List[pycolmap.Image]]: """Map each image to the images it shares at least min_shared_points 3D points with. Returns a dict image_id -> list of overlapping images (excluding self). """ print("Pre-computing overlapping views...") image_3d_points = {} for img_id, image in colmap_reconstruction.images.items(): points_3d = set() for point2D in image.points2D: if point2D.has_point3D(): points_3d.add(point2D.point3D_id) image_3d_points[img_id] = points_3d overlapping_views = {} total_pairs = 0 overlapping_pairs = 0 for img_id_1, image_1 in colmap_reconstruction.images.items(): overlapping_views[img_id_1] = [] points_1 = image_3d_points[img_id_1] if len(points_1) == 0: continue for img_id_2, image_2 in colmap_reconstruction.images.items(): if img_id_1 >= img_id_2: continue total_pairs += 1 points_2 = image_3d_points[img_id_2] shared_points = points_1.intersection(points_2) if len(shared_points) >= min_shared_points: overlapping_pairs += 1 overlapping_views[img_id_1].append(image_2) if img_id_2 not in overlapping_views: overlapping_views[img_id_2] = [] overlapping_views[img_id_2].append(image_1) print(f" Found {overlapping_pairs}/{total_pairs} overlapping pairs") avg_overlaps = np.mean([len(overlaps) for overlaps in overlapping_views.values()]) if overlapping_views else 0 print(f" Average overlaps per image: {avg_overlaps:.1f}") return overlapping_views def check_3d_point_multi_view_consistency( point_3d: np.ndarray, original_vertex_type: str, current_colmap_image: pycolmap.Image, precomputed_overlaps: Dict[int, List[pycolmap.Image]], image_data_map: Dict[str, np.ndarray], gestalt_color_mapping: Dict[str, tuple], min_consistent_views: int = 2, projection_patch_size: int = 5, debug: bool = False ) -> bool: """Check whether a 3D vertex is consistent across the views that see it. Projects the point into each overlapping view and checks that the gestalt segmentation around the projection matches the vertex's class. Returns True if at least min_consistent_views agree, and also True when there are too few overlapping views to verify. """ if original_vertex_type not in gestalt_color_mapping: if debug: print(f" Vertex type {original_vertex_type} not in color mapping") return False target_color = np.array(gestalt_color_mapping[original_vertex_type]) overlapping_views = precomputed_overlaps.get(current_colmap_image.image_id, []) if debug: print(f" Found {len(overlapping_views)} overlapping views for vertex type {original_vertex_type}") if len(overlapping_views) < min_consistent_views: if debug: print(f" Not enough overlapping views ({len(overlapping_views)} < {min_consistent_views}), accepting point") return True # Accept when we cannot verify (too few overlapping views). consistent_view_count = 0 total_checked_views = 0 half_patch = projection_patch_size // 2 for other_image in overlapping_views: try: total_checked_views += 1 projection = other_image.project_point(point_3d) if projection is None: if debug: print(f" View {other_image.name}: projection failed (behind camera)") continue u, v = projection img_width = other_image.camera.width img_height = other_image.camera.height if not (0 <= u < img_width and 0 <= v < img_height): if debug: print(f" View {other_image.name}: projection out of bounds ({u:.1f}, {v:.1f})") continue other_gest_seg_np = None for img_name, gest_seg_np in image_data_map.items(): if img_name in other_image.name or other_image.name in img_name: other_gest_seg_np = gest_seg_np break if other_gest_seg_np is None: if debug: print(f" View {other_image.name}: no segmentation data found") continue seg_h, seg_w = other_gest_seg_np.shape[:2] # Rescale the projection if the segmentation resolution differs from the camera. if seg_w != img_width or seg_h != img_height: u_seg = int(round(u * seg_w / img_width)) v_seg = int(round(v * seg_h / img_height)) else: u_seg, v_seg = int(round(u)), int(round(v)) v_start = max(0, v_seg - half_patch) v_end = min(seg_h, v_seg + half_patch + 1) u_start = max(0, u_seg - half_patch) u_end = min(seg_w, u_seg + half_patch + 1) seg_color_patch = other_gest_seg_np[v_start:v_end, u_start:u_end] if seg_color_patch.size > 0: patch_matches = np.any(np.all(np.abs(seg_color_patch - target_color) <= 1.0, axis=-1)) if patch_matches: consistent_view_count += 1 if debug: print(f" View {other_image.name}: MATCH at ({u:.1f}, {v:.1f}) -> ({u_seg}, {v_seg})") else: if debug: unique_colors = np.unique(seg_color_patch.reshape(-1, 3), axis=0) print(f" View {other_image.name}: no match at ({u:.1f}, {v:.1f}) -> ({u_seg}, {v_seg}), colors: {unique_colors[:3]}") except Exception as e: if debug: print(f" View {other_image.name}: exception {e}") continue result = consistent_view_count >= min_consistent_views if debug: print(f" Result: {consistent_view_count}/{total_checked_views} consistent views, required: {min_consistent_views}, accepted: {result}") return result def filter_vertices_by_multi_view_consistency( vert_edge_per_image: Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]], colmap_reconstruction: pycolmap.Reconstruction, gestalt_segmentations: List[PImage.Image], image_ids: List[str], gestalt_color_mapping: Dict[str, tuple], depth_size_per_image: List[Tuple[int, int]], # [(W, H), ...] for each image min_consistent_views: int = 2, min_shared_points_for_overlap: int = 10, projection_patch_size: int = 25 ) -> Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]]: """Drop 3D vertices that are not multi-view consistent and remap edges to the survivors.""" precomputed_overlaps = precompute_overlapping_views( colmap_reconstruction, min_shared_points_for_overlap ) image_data_map = {} for i, (gest_seg, img_id, (w, h)) in enumerate(zip(gestalt_segmentations, image_ids, depth_size_per_image)): gest_seg_resized = gest_seg.resize((w, h)) gest_seg_np = np.array(gest_seg_resized).astype(np.uint8) image_data_map[img_id] = gest_seg_np filtered_vert_edge_per_image = {} for img_idx, (orig_2d_verts, orig_2d_conns, v3d_candidates) in vert_edge_per_image.items(): if len(v3d_candidates) == 0: filtered_vert_edge_per_image[img_idx] = ([], [], np.empty((0, 3))) continue current_img_id = image_ids[img_idx] current_colmap_img = None for colmap_img_id, colmap_img in colmap_reconstruction.images.items(): if current_img_id in colmap_img.name: current_colmap_img = colmap_img break if current_colmap_img is None: filtered_vert_edge_per_image[img_idx] = (orig_2d_verts, orig_2d_conns, v3d_candidates) continue kept_v3d = [] kept_orig_2d_verts_indices = [] for j, p_3d in enumerate(v3d_candidates): if j >= len(orig_2d_verts): continue original_vertex_type = orig_2d_verts[j]['type'] is_consistent = check_3d_point_multi_view_consistency( p_3d, original_vertex_type, current_colmap_img, precomputed_overlaps, image_data_map, gestalt_color_mapping, min_consistent_views, projection_patch_size=projection_patch_size, debug=False ) if is_consistent: kept_v3d.append(p_3d) kept_orig_2d_verts_indices.append(j) if len(kept_v3d) == 0: filtered_vert_edge_per_image[img_idx] = ([], [], np.empty((0, 3))) continue new_orig_2d_verts = [orig_2d_verts[j] for j in kept_orig_2d_verts_indices] old_idx_to_new_idx = {old_idx: new_idx for new_idx, old_idx in enumerate(kept_orig_2d_verts_indices)} new_orig_2d_conns = [] for (u, v) in orig_2d_conns: if u in old_idx_to_new_idx and v in old_idx_to_new_idx: new_u = old_idx_to_new_idx[u] new_v = old_idx_to_new_idx[v] new_orig_2d_conns.append((new_u, new_v)) filtered_vert_edge_per_image[img_idx] = ( new_orig_2d_verts, new_orig_2d_conns, np.array(kept_v3d) ) return filtered_vert_edge_per_image def recover_edges_after_vertex_filtering( filtered_vert_edge_per_image: Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]], good_entry: dict, edge_class_names: List[str], edge_th: float = 25.0 ) -> Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]]: """Re-detect edges between the surviving vertices after vertex filtering, using the same semantic rulebook and line-of-sight verification as the initial edge detection.""" recovered_vert_edge_per_image = {} total_new_edges = 0 total_original_edges = 0 for img_idx, (filtered_2d_verts, filtered_2d_conns, filtered_v3d) in filtered_vert_edge_per_image.items(): total_original_edges += len(filtered_2d_conns) if len(filtered_2d_verts) < 2: recovered_vert_edge_per_image[img_idx] = (filtered_2d_verts, filtered_2d_conns, filtered_v3d) continue try: gest = good_entry['gestalt'][img_idx] depth = good_entry['depth'][img_idx] depth_size = (np.array(depth).shape[1], np.array(depth).shape[0]) # W, H gest_seg = gest.resize(depth_size) gest_seg_np = np.array(gest_seg).astype(np.uint8) except (IndexError, KeyError): recovered_vert_edge_per_image[img_idx] = (filtered_2d_verts, filtered_2d_conns, filtered_v3d) continue structural_pts = np.array([v['xy'] for v in filtered_2d_verts]) structural_idx_map = list(range(len(filtered_2d_verts))) new_connections = [] for edge_class in edge_class_names: # --- 1. THE SEMANTIC RULEBOOK --- if edge_class in ['ridge']: allowed_types = ['apex'] elif edge_class in ['eave', 'flashing', 'step_flashing']: allowed_types = ['eave_end_point', 'flashing_end_point'] else: # rake, valley, hip, transition_line allowed_types = ['apex', 'eave_end_point', 'flashing_end_point'] allowed_pts = [] allowed_idx_map = [] for orig_idx, v in enumerate(filtered_2d_verts): if v['type'] in allowed_types: allowed_pts.append(v['xy']) allowed_idx_map.append(orig_idx) allowed_pts = np.array(allowed_pts) if len(allowed_pts) < 2: continue # -------------------------------- edge_color = np.array(gestalt_color_mapping[edge_class]) mask_raw = cv2.inRange(gest_seg_np, edge_color-0.5, edge_color+0.5) kernel = np.ones((5, 5), np.uint8) mask = cv2.morphologyEx(mask_raw, cv2.MORPH_CLOSE, kernel) if mask.sum() == 0: continue output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S) (numLabels, labels, stats, centroids) = output stats, centroids = stats[1:], centroids[1:] label_indices = range(1, numLabels) for lbl in label_indices: mask_i = np.zeros_like(mask) mask_i[labels == lbl] = 255 # HoughLinesP for discrete line segments lines = cv2.HoughLinesP(mask_i, rho=1, theta=np.pi/180, threshold=15, minLineLength=8, maxLineGap=20) if lines is None: continue for line in lines: x1, y1, x2, y2 = line[0] p1 = np.array([x1, y1], dtype=np.float32) p2 = np.array([x2, y2], dtype=np.float32) if len(allowed_pts) < 2: continue # Distance check using ONLY the semantically allowed points dists = np.array([ point_to_segment_dist(allowed_pts[i], p1, p2) for i in range(len(allowed_pts)) ]) near_mask = (dists <= edge_th) near_indices = np.where(near_mask)[0] if len(near_indices) < 2: continue # --- 2. CONNECTIVITY WITH LINE-OF-SIGHT VERIFICATION --- 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] vA = allowed_idx_map[idx_a] vB = allowed_idx_map[idx_b] conn = tuple(sorted((vA, vB))) if conn not in new_connections: # THE ULTIMATE SPIDERWEB KILLER: # Verify that the line between these two corners actually exists in the mask! is_valid_edge = verify_edge_mask(allowed_pts[idx_a], allowed_pts[idx_b], mask, min_overlap=0.3) if is_valid_edge: new_connections.append(conn) total_new_edges += len(new_connections) recovered_vert_edge_per_image[img_idx] = (filtered_2d_verts, new_connections, filtered_v3d) print(f" Edge recovery details: {total_original_edges} -> {total_new_edges} edges across all images") return recovered_vert_edge_per_image def merge_collinear_edges(vertices: np.ndarray, edges: List[Tuple[int, int]], cos_threshold: float = -0.98) -> Tuple[np.ndarray, List[Tuple[int, int]]]: """ Finds degree-2 vertices that form a straight line and merges their edges. cos_threshold of -0.98 corresponds to ~170 degrees. """ if len(edges) == 0: return vertices, edges adj = defaultdict(set) for u, v in edges: adj[u].add(v) adj[v].add(u) edges_set = set([tuple(sorted((u, v))) for u, v in edges]) merged_something = True while merged_something: merged_something = False for b in list(adj.keys()): neighbors = list(adj[b]) if len(neighbors) == 2: a, c = neighbors vec1 = vertices[a] - vertices[b] vec2 = vertices[c] - vertices[b] norm1 = np.linalg.norm(vec1) norm2 = np.linalg.norm(vec2) if norm1 > 1e-5 and norm2 > 1e-5: cos_sim = np.dot(vec1, vec2) / (norm1 * norm2) if cos_sim < cos_threshold: e1 = tuple(sorted((a, b))) e2 = tuple(sorted((b, c))) if e1 in edges_set: edges_set.remove(e1) if e2 in edges_set: edges_set.remove(e2) new_edge = tuple(sorted((a, c))) edges_set.add(new_edge) adj[a].remove(b) adj[c].remove(b) adj[a].add(c) adj[c].add(a) del adj[b] merged_something = True break new_edges = list(edges_set) return vertices, new_edges def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]: """Predict the 3D wireframe (vertices and edges) from a dataset entry.""" good_entry = convert_entry_to_human_readable(entry) vert_edge_per_image = {} depth_sizes = [] colmap_rec = good_entry.get('colmap', good_entry.get('colmap_binary')) for i, (gest, depth, K, R, t, img_id, ade_seg) in enumerate(zip(good_entry['gestalt'], good_entry['depth'], good_entry['K'], good_entry['R'], good_entry['t'], good_entry['image_ids'], good_entry['ade'] )): K = np.array(K) R = np.array(R) t = np.array(t) depth_size = (np.array(depth).shape[1], np.array(depth).shape[0]) # W, H depth_sizes.append(depth_size) # resize() can place pixels at the wrong positions; see # https://numpy.org/doc/stable/reference/generated/numpy.ndarray.resize.html gest_seg = gest.resize(depth_size) gest_seg_np = np.array(gest_seg).astype(np.uint8) # Match this image to its COLMAP entry by name (as in get_sparse_depth). found_colmap_img = None for img_id_c, col_img in colmap_rec.images.items(): if img_id in col_img.name: found_colmap_img = col_img break point_class_names = ["apex", "eave_end_point", "flashing_end_point"] edge_class_names = ["eave", "ridge", "rake", "valley", "hip", "flashing", "step_flashing", "transition_line"] vertices, connections = get_vertices_and_edges_from_segmentation( gest_seg_np, point_class_names, edge_class_names, colmap_image=found_colmap_img, colmap_points3D=colmap_rec.points3D, edge_th=25.0, min_3d_points_for_vertex=1, vertex_cluster_eps=25.0, use_colmap_for_vertices=False, patch_size=25 ) if (len(vertices) < 2) or (len(connections) < 1): print(f'Not enough vertices or connections found in image {i}, skipping.') vert_edge_per_image[i] = [], [], np.empty((0, 3)) continue vertices_3d = create_3d_wireframe_single_image( vertices, connections, depth, colmap_rec, img_id, ade_seg ) vert_edge_per_image[i] = vertices, connections, vertices_3d print("Applying multi-view consistency filtering...") total_vertices_before_filtering = sum(len(v3d) for _, _, v3d in vert_edge_per_image.values()) print(f"Total vertices before filtering: {total_vertices_before_filtering}") filtered_vert_edge_per_image = filter_vertices_by_multi_view_consistency( vert_edge_per_image, colmap_rec, good_entry['gestalt'], good_entry['image_ids'], gestalt_color_mapping, depth_sizes, min_consistent_views=1, min_shared_points_for_overlap=3, projection_patch_size=30 ) total_vertices_before = sum(len(v3d) for _, _, v3d in vert_edge_per_image.values()) total_vertices_after = sum(len(v3d) for _, _, v3d in filtered_vert_edge_per_image.values()) print(f"Multi-view filtering: {total_vertices_before} -> {total_vertices_after} vertices") print(f"Filtering removed {total_vertices_before - total_vertices_after} vertices ({100*(total_vertices_before - total_vertices_after)/max(total_vertices_before,1):.1f}%)") print("Recovering edges between filtered vertices...") edges_before_recovery = sum(len(conns) for _, conns, _ in filtered_vert_edge_per_image.values()) edge_class_names = ["eave", "ridge", "rake", "valley", "hip", "flashing", "step_flashing", "transition_line"] recovered_vert_edge_per_image = recover_edges_after_vertex_filtering( filtered_vert_edge_per_image, good_entry, edge_class_names, edge_th=25.0 ) edges_after_recovery = sum(len(conns) for _, conns, _ in recovered_vert_edge_per_image.values()) print(f"Edge recovery: {edges_before_recovery} -> {edges_after_recovery} edges") all_3d_vertices, connections_3d = merge_vertices_3d(recovered_vert_edge_per_image, point_class_names, 0.7) all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d, keep_largest=False) # Merge fragmented collinear segments into continuous edges, then drop any # vertices orphaned by the merge. all_3d_vertices_clean, connections_3d_clean = merge_collinear_edges(all_3d_vertices_clean, connections_3d_clean, cos_threshold=-0.98) all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices_clean, connections_3d_clean, keep_largest=False) if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1: print (f'Not enough vertices or connections in the 3D vertices') return empty_solution() return all_3d_vertices_clean, connections_3d_clean