solution.py

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