handcrafted_solution.py

# Description: This file contains the handcrafted solution for the task of wireframe reconstruction 

import io
from PIL import Image as PImage
import numpy as np
from collections import defaultdict
import cv2
from typing import Tuple, List
from scipy.spatial.distance import cdist
from scipy.optimize import minimize



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):
    import pycolmap
    import tempfile,zipfile
    import io
    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']:
            out[k] = np.array(v)
        else:
            out[k]=v
    out['__key__'] = entry['order_id']
    return out


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.
    """
    from hoho.color_mappings import gestalt_color_mapping  # for apex, eave_end_point, etc.

    #--------------------------------------------------------------------------------
    # Step A: Collect apex and eave_end vertices
    #--------------------------------------------------------------------------------
    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, depth_fitted, sparse_depth, search_radius=10):
    """
    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 median (robust) as the vertex 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([v['xy'] for v 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)

    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)
        
        region = sparse_depth[y0:y1, x0:x1]
        valid_vals = region[region > 0]
        if len(valid_vals) > 0:
            # Use median of valid sparse depth
            vertex_depth[i] = np.median(valid_vals)
        else:
            # Fallback to depth_fitted at this pixel
            vertex_depth[i] = depth_fitted[y_i, x_i]
    
    return uv, vertex_depth

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'''
    all_3d_vertices = []
    connections_3d = []
    all_indexes = []
    cur_start = 0
    types = []
    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)
    #print (connections_3d)
    distmat = cdist(all_3d_vertices, all_3d_vertices)
    types = np.array(types).reshape(-1,1)
    same_types = cdist(types, types)
    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])))
    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
    #print (connections_3d)
    new_vertices=np.array(new_vertices)
    #print (connections_3d)
    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)
    #print (f'{len(new_vertices)} left after merging {len(all_3d_vertices)} with {th=}')
    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, K, R, t, depth):
    H, W = depth.shape
    xyz = []
    rgb = []
    found = False
    #print (img_id)
    for img_id_c, col_img in colmap_rec.images.items():
        print (col_img.name)
        if col_img.name == img_id:
            found = True
            break
    if not found:
        print (f"{img_id} not found, returning empty depth")
        return np.zeros((H, W), dtype=np.float32), False
    mat4x4 = np.eye(4)
    mat4x4[:3 ] = col_img.cam_from_world.matrix()
    for pid,p in colmap_rec.points3D.items():
        if col_img.has_point3D(pid):
            xyz.append(p.xyz)
            rgb.append(p.color)
    xyz = np.array(xyz)
    rgb = np.array(rgb)
    xyz_projected = cv2.transform(cv2.convertPointsToHomogeneous(xyz), mat4x4)
    xyz_projected = cv2.convertPointsFromHomogeneous(xyz_projected).reshape(-1, 3)
    uv, _ = cv2.projectPoints(xyz_projected, np.zeros(3), np.zeros(3), np.array(K), np.zeros(4))
    uv = uv.squeeze()
    u, v = uv[:, 0].astype(np.int32), uv[:, 1].astype(np.int32)
    mask = (u >= 0) & (u < W) & (v >= 0) & (v < H)
    u, v = u[mask], v[mask]
    xyz_projected, rgb = xyz_projected[mask], rgb[mask]
    depth = np.zeros((H, W), dtype=np.float32)
    depth[v, u] = xyz_projected[:, 2]
    return depth, True


def fit_scale_robust_median(depth, sparse_depth):
    """
    Fits the model sparse_depth ~ k * depth + b by minimizing the median
    of absolute residuals, i.e. median( |sparse_depth - (k*depth + b)| ).

    Parameters
    ----------
    depth : np.ndarray
        Array of depth estimates (same shape as sparse_depth).
    sparse_depth : np.ndarray
        Sparse array with precise depth at certain locations
        (0 where data is unavailable).

    Returns
    -------
    k : float
        The slope of the robust best-fit affine transform.
    b : float
        The intercept of the robust best-fit affine transform.
    depth_fitted : np.ndarray
        The depth array adjusted by the affine fit: k*depth + b.
    """

    # 1. Create mask of valid (nonzero) locations in sparse_depth
    mask = (sparse_depth != 0)
    X = depth[mask]
    Y = sparse_depth[mask]

    # 2. Define the objective: median of absolute residuals
    def median_abs_resid(params, xvals, yvals):
        k, b = params
        return np.median(np.abs(yvals - (k*xvals)))

    k_init, b_init = 1.0, 0.0
    # 4. Optimize using a derivative-free method (Nelder-Mead)
    res = minimize(
        fun=median_abs_resid,
        x0=[k_init, b_init],
        args=(X, Y),
        method='Nelder-Mead'
    )

    k_robust, b_robust = res.x

    # 5. Construct the fitted depth array
    depth_fitted = k_robust * depth #+ b_robust

    return k_robust, depth_fitted




def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
    good_entry = convert_entry_to_human_readable(entry)
    vert_edge_per_image = {}
    for i, (gest, depth, K, R, t, img_id) in enumerate(zip(good_entry['gestalt'],
                                                good_entry['depth'], 
                                                good_entry['K'],
                                                good_entry['R'],
                                                good_entry['t'],
                                                good_entry['image_ids']
                                                )):
        colmap_rec = good_entry['colmap_binary']
        K = np.array(K)
        R = np.array(R)
        t = np.array(t)
        gest_seg = gest.resize(depth.size)
        gest_seg_np = np.array(gest_seg).astype(np.uint8)
        # Metric3D
        depth_np = np.array(depth) / 1000.
        depth_sparse, found = get_sparse_depth(colmap_rec, img_id, K, R, t, depth_np)
        if not found:
            print (f'No sparse depth found for image {i}')
            vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
            continue
        k, depth_fitted = fit_scale_robust_median(depth_np, depth_sparse)#fit_affine_robust_median(depth_np, depth_sparse)
        print (k)
        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 50.)
        if (len(vertices) < 2) or (len(connections) < 1):
            print (f'Not enough vertices or connections in image {i}')
            vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
            continue

        uv, depth_vert = get_uv_depth(vertices, depth_fitted, depth_sparse, 50)
        # Normalize the uv to the camera intrinsics
        X = (uv[:, 0] - K[0, 2]) / K[0, 0] * depth_vert
        Y = (uv[:, 1] - K[1, 2]) / K[1, 1] * depth_vert
        Z = depth_vert
        vertices_3d_local = np.column_stack([X, Y, Z])
        world_to_cam = np.eye(4)
        world_to_cam[:3, :3] = R
        world_to_cam[:3, 3] = t.reshape(-1)
        cam_to_world =  np.linalg.inv(world_to_cam)
        vertices_3d = cv2.transform(cv2.convertPointsToHomogeneous(vertices_3d_local), cam_to_world)
        vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
        vert_edge_per_image[i] = vertices, connections, vertices_3d
    all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 0.25)
    all_3d_vertices_clean, connections_3d_clean  = prune_not_connected(all_3d_vertices, connections_3d, 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()
    if visualize:
        from hoho.viz3d import plot_estimate_and_gt
        plot_estimate_and_gt(   all_3d_vertices_clean, 
                                connections_3d_clean, 
                                good_entry['wf_vertices'],
                                good_entry['wf_edges'])
    return all_3d_vertices_clean, connections_3d_clean