hoho2025/example_solutions.py

# Description: This file contains the handcrafted solution for the task of wireframe reconstruction 
import io
import tempfile
import zipfile
from collections import defaultdict
from typing import Tuple, List
import cv2
import numpy as np
import pycolmap
from PIL import Image as PImage
from scipy.spatial.distance import cdist

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)  # unpacks cameras.txt, images.txt, etc. to tmpdir
        # Now parse with pycolmap
        rec = pycolmap.Reconstruction(tmpdir)
        return rec

def convert_entry_to_human_readable(entry):
    out = {}
    for k, v in entry.items():
        if 'colmap' in k:
            out[k] = read_colmap_rec(v)
        elif k in ['wf_vertices', 'wf_edges', 'K', 'R', 't', 'depth']:
            out[k] = np.array(v)
        else:
            out[k]=v
    out['__key__'] = entry['order_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 both endpoints are the same, just return distance to one of them
    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 get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=25.0):
    """
    Identify apex and eave-end vertices, then detect lines for eave/ridge/rake/valley.
    For each connected component, we do a line fit with cv2.fitLine, then measure
    segment endpoints more robustly. We then associate apex points that are within
    'edge_th' of the line segment. We record those apex–apex connections for edges
    if at least 2 apexes lie near the same component line.
    """
    #--------------------------------------------------------------------------------
    # Step A: Collect apex and eave_end vertices
    #--------------------------------------------------------------------------------
    if not isinstance(gest_seg_np, np.ndarray):
        gest_seg_np = np.array(gest_seg_np)
    vertices = []
    # Apex
    apex_color = np.array(gestalt_color_mapping['apex'])
    apex_mask = cv2.inRange(gest_seg_np, apex_color-0.5, apex_color+0.5)
    if apex_mask.sum() > 0:
        output = cv2.connectedComponentsWithStats(apex_mask, 8, cv2.CV_32S)
        (numLabels, labels, stats, centroids) = output
        stats, centroids = stats[1:], centroids[1:]  # skip background
        for i in range(numLabels-1):
            vert = {"xy": centroids[i], "type": "apex"}
            vertices.append(vert)

    # Eave end
    eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
    eave_end_mask = cv2.inRange(gest_seg_np, eave_end_color-0.5, eave_end_color+0.5)
    if eave_end_mask.sum() > 0:
        output = cv2.connectedComponentsWithStats(eave_end_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": "eave_end_point"}
            vertices.append(vert)

    # Consolidate apex points as array:
    apex_pts = []
    apex_idx_map = []  # keep track of index in 'vertices'
    for idx, v in enumerate(vertices):
        apex_pts.append(v['xy'])
        apex_idx_map.append(idx)
    apex_pts = np.array(apex_pts)

    connections = []
    edge_classes = ['eave', 'ridge', 'rake', 'valley']
    for edge_class in edge_classes:
        edge_color = np.array(gestalt_color_mapping[edge_class])
        mask_raw = cv2.inRange(gest_seg_np, edge_color-0.5, edge_color+0.5)
        # Possibly do morphological open/close to avoid merges or small holes
        kernel = np.ones((5, 5), np.uint8)  # smaller kernel to reduce over-merge
        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)
            # (vx, vy, x0, y0) = direction + a point on the line
            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()
            # We'll approximate endpoints by projecting (xs, ys) onto the line,
            # then taking min and max in the 1D param along the line.

            # param along the line = ( (x - x0)*vx + (y - y0)*vy )
            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])

            #--------------------------------------------------------------------------------
            # Step C: If apex points are within 'edge_th' of segment, they are connected
            #--------------------------------------------------------------------------------
            if len(apex_pts) < 2:
                continue

            # Distance from each apex to the line segment
            dists = np.array([
                point_to_segment_dist(apex_pts[i], p1, p2)
                for i in range(len(apex_pts))
            ])

            # Indices of apex points that are near
            near_mask = (dists <= edge_th)
            near_indices = np.where(near_mask)[0]
            if len(near_indices) < 2:
                continue

            # Connect each pair among these near apex points
            for i in range(len(near_indices)):
                for j in range(i+1, len(near_indices)):
                    a_idx = near_indices[i]
                    b_idx = near_indices[j]
                    # 'a_idx' and 'b_idx' are indices in apex_pts / apex_idx_map
                    vA = apex_idx_map[a_idx]
                    vB = apex_idx_map[b_idx]
                    # Store the connection using sorted indexing
                    conn = tuple(sorted((vA, vB)))
                    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, returns a 2D array of (u,v) and a matching 1D array of depths.
    
    We attempt to use the sparse_depth if available in a local neighborhood:
      1. For each vertex coordinate (x, y), define a local window in sparse_depth 
         of size (2*search_radius + 1).
      2. Collect all valid (nonzero) values in that window.
      3. If any exist, we take the *closest* valid pixel's depth.
      4. Otherwise, we use depth_fitted[y, x].
    
    Parameters
    ----------
    vertices : List[dict]
        Each dict must have "xy" at least, e.g. {"xy": (x, y), ...}
    depth_fitted : np.ndarray
        A 2D array (H, W), the dense (or corrected) depth for fallback.
    sparse_depth : np.ndarray
        A 2D array (H, W), mostly zeros except where accurate data is available.
    search_radius : int
        Pixel radius around the vertex in which to look for sparse depth values.
    
    Returns
    -------
    uv : np.ndarray of shape (N, 2)
        2D float coordinates of each vertex (x, y).
    vertex_depth : np.ndarray of shape (N,)
        Depth value chosen for each vertex.
    """
    
    # Collect each vertex's (x, y)
    uv = np.array([vert['xy'] for vert in vertices], dtype=np.float32)
    
    # Convert to integer pixel coordinates (round or floor)
    uv_int = np.round(uv).astype(np.int32)
    H, W = depth_fitted.shape[:2]
    
    # Clip coordinates to stay within image bounds
    uv_int[:, 0] = np.clip(uv_int[:, 0], 0, W - 1)
    uv_int[:, 1] = np.clip(uv_int[:, 1], 0, H - 1)
    
    # Prepare output array of depths
    vertex_depth = np.zeros(len(vertices), dtype=np.float32)
    dense_count = 0
    
    for i, (x_i, y_i) in enumerate(uv_int):
        # Local region in [x_i - search_radius, x_i + search_radius]
        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)
        
        # Crop out the local window in sparse_depth
        region = sparse_depth[y0:y1, x0:x1]
        
        # Find all valid (non-zero) depths
        valid_mask = (region > 0)
        valid_y, valid_x = np.where(valid_mask)
        
        if valid_y.size > 0:
            # Compute global coordinates for each valid pixel
            global_x = x0 + valid_x
            global_y = y0 + valid_y
            
            # Compute squared distance to center (x_i, y_i)
            dist_sq = (global_x - x_i)**2 + (global_y - y_i)**2
            
            # Find the nearest valid pixel
            min_idx = np.argmin(dist_sq)
            nearest_depth = region[valid_y[min_idx], valid_x[min_idx]]
            vertex_depth[i] = nearest_depth
        else:
            # Fallback to the dense depth
            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) -> np.ndarray:
    """
    Projects 2D vertex coordinates with associated depths to 3D world coordinates.

    Parameters
    ----------
    uv : np.ndarray
        (N, 2) array of 2D vertex coordinates (u, v).
    depth_vert : np.ndarray
        (N,) array of depth values for each vertex.
    col_img : pycolmap.Image

    Returns
    -------
    vertices_3d : np.ndarray
        (N, 3) array of vertex coordinates in 3D world space.
    """
    # Backproject to 3D local camera coordinates
    xy_local = np.ones((len(uv), 3))
    K = col_img.camera.calibration_matrix()
    xy_local[:, 0] = (uv[:, 0] - K[0, 2]) / K[0, 0]
    xy_local[:, 1] = (uv[:, 1] - K[1, 2]) / K[1, 1]
    # Get the 3D vertices
    vertices_3d_local = xy_local * depth_vert[...,None]
    
    # Create camera-to-world transformation matrix
    world_to_cam = np.eye(4)
    world_to_cam[:3] = col_img.cam_from_world.matrix()
    cam_to_world = np.linalg.inv(world_to_cam)
    
    # Transform local 3D points to world coordinates
    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,
                                     colmap_rec: pycolmap.Reconstruction,
                                     img_id: str,
                                     ade_seg: PImage) -> np.ndarray:
    """
    Processes a single image view to generate 3D vertex coordinates from existing 2D vertices/edges.

    Parameters
    ----------
    vertices : List[dict]
        List of 2D vertex dictionaries (e.g., {"xy": (x, y), "type": ...}).
    connections : List[Tuple[int, int]]
        List of 2D edge connections (indices into the vertices list).
    depth : PIL.Image
        Initial dense depth map as a PIL Image.
    colmap_rec : pycolmap.Reconstruction
        COLMAP reconstruction data.
    img_id : str
        Identifier for the current image within the COLMAP reconstruction.
    ade_seg : PIL.Image
        ADE20k segmentation map for the image.

    Returns
    -------
    vertices_3d : np.ndarray
        (N, 3) array of vertex coordinates in 3D world space.
        Returns an empty array if processing fails (e.g., missing sparse depth).
    """
    # Check if initial vertices/connections are valid
    if (len(vertices) < 2) or (len(connections) < 1):
        # This case should ideally be handled before calling, but good to double check.
        print(f'Warning: create_3d_wireframe_single_image called with insufficient vertices/connections for image {img_id}')
        return np.empty((0, 3))

    # Get fitted dense depth and sparse depth
    depth_fitted, depth_sparse, found_sparse, col_img = get_fitted_dense_depth(
        depth, colmap_rec, img_id, ade_seg
    )

    # Get UV coordinates and depth for each vertex
    uv, depth_vert = get_uv_depth(vertices, depth_fitted, depth_sparse, 10)

    # Backproject to 3D
    vertices_3d = project_vertices_to_3d(uv, depth_vert, col_img)

    return vertices_3d


def merge_vertices_3d(vert_edge_per_image, th=0.5):
    '''Merge vertices that are close to each other in 3D space and are of same types'''
    # Initialize structures to collect vertices and connections from all images
    all_3d_vertices = []
    connections_3d = []
    all_indexes = []
    cur_start = 0
    types = []
    
    # Combine vertices and update connection indices across all images
    for cimg_idx, (vertices, connections, vertices_3d) in vert_edge_per_image.items():
        types += [int(v['type']=='apex') for v in vertices]
        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)
    
    # Calculate distance matrix between all vertices
    distmat = cdist(all_3d_vertices, all_3d_vertices)
    types = np.array(types).reshape(-1,1)
    same_types = cdist(types, types)
    
    # Create mask for vertices that should be merged (close in space and same type)
    mask_to_merge = (distmat <= th) & (same_types==0)
    new_vertices = []
    new_connections = []
    
    # Extract vertex indices to merge based on the mask
    to_merge = sorted(list(set([tuple(a.nonzero()[0].tolist()) for a in mask_to_merge])))
    
    # Build groups of vertices to merge (transitive grouping)
    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
    
    # Remove duplicates in each group
    for k, v in to_merge_final.items():
        to_merge_final[k] = list(set(v))
    
    # Create final merge groups without duplicates
    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)
    
    # Calculate new vertex positions (average of merged groups)
    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)
    
    # Update connections to use new vertex indices
    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.array([]), []

    # adjacency
    adj = defaultdict(set)
    for (i, j) in connections_3d:
        adj[i].add(j)
        adj[j].add(i)

    # keep only vertices that appear in at least one edge
    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

    # Collect all subgraphs
    comps = []
    for idx in used_idxs:
        if idx not in visited:
            c = bfs(idx)
            comps.append(c)

    # pick largest
    comps.sort(key=lambda c: len(c), reverse=True)
    largest = comps[0] if len(comps)>0 else []

    # Remap
    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]))

    # remove duplicates
    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):
    """
    Return a sparse depth map for the COLMAP image whose name contains
    `img_id_substring`. The output is an array of shape `depth_shape` (H,W),
    where only the projected 3D points get a depth > 0, else 0.
    """
    H, W = depth.shape

    # 1) Find the matching COLMAP image
    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:
        print(f"Image substring {img_id_substring} not found in COLMAP.")
        return np.zeros((H, W), dtype=np.float32), False, None
    
    # 2) Gather 3D points that this image sees
    points_xyz = []
    for pid, p3D in colmap_rec.points3D.items():
        if found_img.has_point3D(pid):
            points_xyz.append(p3D.xyz)  # world coords
    if not points_xyz:
        print(f"No 3D points associated with {found_img.name}.")
        return np.zeros((H, W), dtype=np.float32), False, found_img
    
    points_xyz = np.array(points_xyz)  # (N, 3)
    
    # 3) For each point, project via col_img.project_point()
    uv = []
    z_vals = []
    for xyz in points_xyz:
        proj = found_img.project_point(xyz)  # returns (u, v) in image coords or None
        if proj is not None:
            u_i, v_i = proj
            u_i = int(round(u_i))
            v_i = int(round(v_i))
            # Check in-bounds
            if 0 <= u_i < W and 0 <= v_i < H:
                uv.append((u_i, v_i))
                # We'll compute depth as Z in camera coords
                # from the world->cam transform col_img holds
                mat4x4 = np.eye(4)
                mat4x4[:3, :4] = found_img.cam_from_world.matrix()
                p_cam =  mat4x4@ np.array([xyz[0], xyz[1], xyz[2], 1.0])
                z_vals.append(p_cam[2] / p_cam[3]) 
    
    uv = np.array(uv, dtype=int)     # shape (M,2)
    z_vals = np.array(z_vals)        # shape (M,)
    
    depth_out = np.zeros((H, W), dtype=np.float32)
    depth_out[uv[:,1], uv[:,0]] = z_vals  # Note: uv = (u, v), so row = v, col = u
    
    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):
    """
    Gets sparse depth from COLMAP, computes a house mask, fits dense depth to sparse 
    depth within the mask, and returns the fitted dense depth.

    Parameters
    ----------
    depth : np.ndarray
        Initial dense depth map (H, W).
    colmap_rec : pycolmap.Reconstruction
        COLMAP reconstruction data.
    img_id : str
        Identifier for the current image within the COLMAP reconstruction.
    K : np.ndarray
        Camera intrinsic matrix (3x3).
    R : np.ndarray
        Camera rotation matrix (3x3).
    t : np.ndarray
        Camera translation vector (3,).
    ade20k_seg : PIL.Image
        ADE20k segmentation map for the image.

    Returns
    -------
    depth_fitted : np.ndarray
        Dense depth map scaled and shifted to align with sparse depth within the house mask (H, W).
    depth_sparse : np.ndarray
        The sparse depth map obtained from COLMAP (H, W).
    found_sparse : bool
        True if sparse depth points were found for this image, False otherwise.
    """
    depth_np = np.array(depth) / 1000. # Convert mm to meters if needed
    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 original (meter-scaled) depth if no sparse data
        return depth_np, np.zeros_like(depth_np), False, None

    # Get house mask to focus fitting on relevant areas
    house_mask = get_house_mask(ade20k_seg)
    
    # Fit dense depth to sparse depth (scale only), using only points within the house mask
    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}")
    #depth_fitted = depth_np# * house_mask.astype(np.float32)
    depth_sparse = depth_sparse# * house_mask.astype(np.float32)
    return depth_fitted, depth_sparse, True, col_img


def prune_too_far(all_3d_vertices, connections_3d, colmap_rec, th = 3.0):
    """
    Prune vertices that are too far from sparse point cloud
    
    """
    xyz_sfm=[]
    for k, v in colmap_rec.points3D.items():
        xyz_sfm.append(v.xyz)
    xyz_sfm = np.array(xyz_sfm)
    distmat = cdist(all_3d_vertices, xyz_sfm)
    mindist = distmat.min(axis=1)
    mask = mindist <= th
    all_3d_vertices_new = all_3d_vertices[mask]
    old_idx_survived = np.arange(len(all_3d_vertices))[mask]
    new_idxs = np.arange(len(all_3d_vertices_new))
    old_to_new_idx = dict(zip(old_idx_survived, new_idxs))
    connections_3d_new = [(old_to_new_idx[conn[0]], old_to_new_idx[conn[1]]) for conn in connections_3d if mask[conn[0]] and mask[conn[1]]]   
    return all_3d_vertices_new, connections_3d_new


def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
    """
    Predict 3D wireframe from a dataset entry.
    """
    good_entry = convert_entry_to_human_readable(entry)
    vert_edge_per_image = {}
    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'] # Added ade20k segmentation
                                                )):
        colmap_rec = good_entry['colmap_binary']
        K = np.array(K)
        R = np.array(R)
        t = np.array(t)
        # Resize gestalt segmentation to match depth map size
        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)
        
        # Get 2D vertices and edges first
        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=10.)
        
        # Check if we have enough to proceed
        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
            
        # Call the refactored function to get 3D points
        vertices_3d = create_3d_wireframe_single_image(
            vertices, connections, depth, colmap_rec, img_id, ade_seg
        )
        # Store original 2D vertices, connections, and computed 3D points
        vert_edge_per_image[i] = vertices, connections, vertices_3d
    
    # 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)
    all_3d_vertices_clean, connections_3d_clean  = prune_too_far(all_3d_vertices_clean, connections_3d_clean, colmap_rec, th = 4.0)
    
    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