# Description: This file contains the handcrafted solution for the task of wireframe reconstruction import io from collections import defaultdict from typing import Tuple, List import cv2 import hoho import numpy as np import scipy.interpolate as si from PIL import Image as PImage from hoho.color_mappings import gestalt_color_mapping from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary from scipy.spatial import KDTree from scipy.spatial.distance import cdist import scipy.cluster.hierarchy as shc apex_color = gestalt_color_mapping["apex"] eave_end_point = gestalt_color_mapping["eave_end_point"] flashing_end_point = gestalt_color_mapping["flashing_end_point"] apex_color, eave_end_point, flashing_end_point = [np.array(i) for i in [apex_color, eave_end_point, flashing_end_point]] unclassified = np.array([(215, 62, 138)]) line_classes = ['eave', 'ridge', 'rake', 'valley'] def empty_solution(): '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.''' return np.zeros((2, 3)), [(0, 1)] def convert_entry_to_human_readable(entry): out = {} already_good = {'__key__', 'wf_vertices', 'wf_edges', 'edge_semantics', 'mesh_vertices', 'mesh_faces', 'face_semantics', 'K', 'R', 't'} for k, v in entry.items(): if k in already_good: out[k] = v continue match k: case 'points3d': out[k] = read_points3D_binary(fid=io.BytesIO(v)) case 'cameras': out[k] = read_cameras_binary(fid=io.BytesIO(v)) case 'images': out[k] = read_images_binary(fid=io.BytesIO(v)) case 'ade20k' | 'gestalt': out[k] = [PImage.open(io.BytesIO(x)).convert('RGB') for x in v] case 'depthcm': out[k] = [PImage.open(io.BytesIO(x)) for x in entry['depthcm']] return out def remove_undesired_objects(image): image = image.astype('uint8') nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(image, connectivity=4) sizes = stats[:, -1] max_label = 1 max_size = sizes[1] for i in range(2, nb_components): if sizes[i] > max_size: max_label = i max_size = sizes[i] img2 = np.zeros(output.shape) img2[output == max_label] = 1 return img2 def clean_image(image_gestalt) -> np.ndarray: # clears image in from of unclassified and disconected components image_gestalt = np.array(image_gestalt) unclassified_mask = cv2.inRange(image_gestalt, unclassified + 0.0, unclassified + 0.8) unclassified_mask = cv2.bitwise_not(unclassified_mask) mask = remove_undesired_objects(unclassified_mask).astype(np.uint8) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((11, 11), np.uint8), iterations=11) mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, np.ones((11, 11), np.uint8), iterations=2) image_gestalt[:, :, 0] *= mask image_gestalt[:, :, 1] *= mask image_gestalt[:, :, 2] *= mask return image_gestalt def get_vertices(image_gestalt, *, color_range=4., dialations=3, erosions=1, kernel_size=13): ### detects the apex and eave end and flashing end points apex_mask = cv2.inRange(image_gestalt, apex_color - color_range, apex_color + color_range) eave_end_point_mask = cv2.inRange(image_gestalt, eave_end_point - color_range, eave_end_point + color_range) flashing_end_point_mask = cv2.inRange(image_gestalt, flashing_end_point - color_range, flashing_end_point + color_range) eave_end_point_mask = cv2.bitwise_or(eave_end_point_mask, flashing_end_point_mask) kernel = np.ones((kernel_size, kernel_size), np.uint8) apex_mask = cv2.morphologyEx(apex_mask, cv2.MORPH_DILATE, kernel, iterations=dialations) apex_mask = cv2.morphologyEx(apex_mask, cv2.MORPH_ERODE, kernel, iterations=erosions) eave_end_point_mask = cv2.morphologyEx(eave_end_point_mask, cv2.MORPH_DILATE, kernel, iterations=dialations) eave_end_point_mask = cv2.morphologyEx(eave_end_point_mask, cv2.MORPH_ERODE, kernel, iterations=erosions) *_, apex_centroids = cv2.connectedComponentsWithStats(apex_mask, connectivity=4, stats=cv2.CV_32S) *_, other_centroids = cv2.connectedComponentsWithStats(eave_end_point_mask, connectivity=4, stats=cv2.CV_32S) return apex_centroids[1:], other_centroids[1:], apex_mask, eave_end_point_mask def infer_vertices(image_gestalt, *, color_range=4.): ridge_color = np.array(gestalt_color_mapping["ridge"]) rake_color = np.array(gestalt_color_mapping["rake"]) ridge_mask = cv2.inRange(image_gestalt, ridge_color - color_range, ridge_color + color_range) ridge_mask = cv2.morphologyEx(ridge_mask, cv2.MORPH_DILATE, np.ones((3, 3)), iterations=4) rake_mask = cv2.inRange(image_gestalt, rake_color - color_range, rake_color + color_range) rake_mask = cv2.morphologyEx(rake_mask, cv2.MORPH_DILATE, np.ones((3, 3)), iterations=4) intersection_mask = cv2.bitwise_and(ridge_mask, rake_mask) intersection_mask = cv2.morphologyEx(intersection_mask, cv2.MORPH_DILATE, np.ones((11, 11)), iterations=3) *_, inferred_centroids = cv2.connectedComponentsWithStats(intersection_mask, connectivity=4, stats=cv2.CV_32S) return inferred_centroids[1:], intersection_mask def get_missed_vertices(vertices, inferred_centroids, *, min_missing_distance=200.0, **kwargs): vertices = KDTree(vertices) closest = vertices.query(inferred_centroids, k=1, distance_upper_bound=min_missing_distance) missed_points = inferred_centroids[closest[1] == len(vertices.data)] return missed_points def get_lines_and_directions(gest_seg_np, edge_class, *, color_range=4., rho, theta, threshold, min_line_length, max_line_gap, extend, **kwargs): edge_color = np.array(gestalt_color_mapping[edge_class]) mask = cv2.inRange(gest_seg_np, edge_color - color_range, edge_color + color_range) mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, np.ones((3, 3)), iterations=1) if not np.any(mask): return [], [] # Run Hough on edge detected image # Output "lines" is an array containing endpoints of detected line segments cv2.GaussianBlur(mask, (11, 11), 0, mask) lines = cv2.HoughLinesP(mask, rho, theta, threshold, np.array([]), min_line_length, max_line_gap) if lines is None: return [], [] line_directions = [] edges = [] for line_idx, line in enumerate(lines): for x1, y1, x2, y2 in line: if x1 < x2: x1, y1, x2, y2 = x2, y2, x1, y1 direction = (np.array([x2 - x1, y2 - y1])) direction = direction / np.linalg.norm(direction) line_directions.append(direction) direction = extend * direction x1, y1 = (-direction + (x1, y1)).astype(np.int32) x2, y2 = (+ direction + (x2, y2)).astype(np.int32) edges.append((x1, y1, x2, y2)) return edges, line_directions def infer_missing_vertices(ridge_edges, rake_edges): ridge_edges = np.array(ridge_edges) rake_edges = np.array(rake_edges) ridge_ends = np.concatenate([ridge_edges[:, 2:], ridge_edges[:, :2]]) rake_ends = np.concatenate([rake_edges[:, 2:], rake_edges[:, :2]]) ridge_ends = KDTree(ridge_ends) rake_ends = KDTree(rake_ends) missing_candidates = rake_ends.query_ball_tree(ridge_ends, 5) missing_candidates = np.concatenate([*missing_candidates]) missing_candidates = np.unique(missing_candidates).astype(np.int32) return ridge_ends.data[missing_candidates] def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, max_angle=5., **kwargs): '''Get the vertices and edges from the gestalt segmentation mask of the house''' # Apex connections = [] deviation_threshold = np.cos(np.deg2rad(max_angle)) apex_centroids, eave_end_point_centroids, apex_mask, eave_end_point_mask = get_vertices(gest_seg_np) vertices = np.concatenate([apex_centroids, eave_end_point_centroids]) # inferred_vertices, inferred_mask = infer_vertices(gest_seg_np) # missed_vertices = get_missed_vertices(vertices, inferred_vertices, **kwargs) # vertices = np.concatenate([vertices, missed_vertices]) if len(vertices) < 2: return [], [] # scale = 1 # vertex_size = np.zeros(vertices.shape[0]) # for i, coords in enumerate(vertices): # # coords = np.round(coords).astype(np.uint32) # radius = point_radius # np.clip(int(max_depth//2 + depth_np[coords[1], coords[0]]), 10, 30)#int(np.clip(max_depth - depth_np[coords[1], coords[0]], 10, 20)) # vertex_size[i] = (scale * radius) ** 2 # because we are using squared distances edges = [] line_directions = [] rho = 1 # distance resolution in pixels of the Hough grid theta = np.pi / 180 # angular resolution in radians of the Hough grid threshold = 20 # minimum number of votes (intersections in Hough grid cell) min_line_length = 60 # minimum number of pixels making up a line max_line_gap = 40 # maximum gap in pixels between connectable line segments ridge_edges, ridge_directions = get_lines_and_directions(gest_seg_np, "ridge", rho=rho, theta=theta, threshold=threshold, min_line_length=min_line_length, max_line_gap=max_line_gap, **kwargs) rake_edges, rake_directions = get_lines_and_directions(gest_seg_np, "rake", rho=rho, theta=theta, threshold=threshold, min_line_length=min_line_length, max_line_gap=max_line_gap, **kwargs) if len(ridge_edges) > 0: edges.append(ridge_edges) line_directions.append(ridge_directions) if len(rake_edges) > 0: edges.append(rake_edges) line_directions.append(rake_directions) missed_vertices = [] if len(ridge_edges) > 0 and len(rake_edges) > 0: inferred_vertices = infer_missing_vertices(ridge_edges, rake_edges) missed_vertices = get_missed_vertices(vertices, inferred_vertices, **kwargs) vertices = np.concatenate([vertices, missed_vertices]) vertices = KDTree(vertices) for edge_class in ['eave', 'step_flashing', 'flashing', 'post', 'valley', 'hip', 'transition_line']: class_edges, class_directions = get_lines_and_directions(gest_seg_np, edge_class, rho=rho, theta=theta, threshold=threshold, min_line_length=min_line_length, max_line_gap=max_line_gap, **kwargs) if len(class_edges) > 0: edges.append(class_edges) line_directions.append(class_directions) edges = np.concatenate(edges).astype(np.float64) line_directions = np.concatenate(line_directions).astype(np.float64) if len(edges) < 1: return [], [] # calculate the distances between the vertices and the edge ends begin_edges = KDTree(edges[:, :2]) end_edges = KDTree(edges[:, 2:]) begin_indices = begin_edges.query_ball_tree(vertices, point_radius) end_indices = end_edges.query_ball_tree(vertices, point_radius) line_indices = np.where(np.array([len(i) and len(j) for i, j in zip(begin_indices, end_indices)]))[0] # create all possible connections between begin and end candidates that correspond to a line begin_vertex_list = [] end_vertex_list = [] line_idx_list = [] for line_idx in line_indices: begin_vertex, end_vertex = begin_indices[line_idx], end_indices[line_idx] begin_vertex, end_vertex = np.meshgrid(begin_vertex, end_vertex) begin_vertex_list.extend(begin_vertex.flatten()) end_vertex_list.extend(end_vertex.flatten()) line_idx_list.extend([line_idx] * len(begin_vertex.flatten())) line_idx_list = np.array(line_idx_list) all_connections = np.array([begin_vertex_list, end_vertex_list]) # decrease the number of possible connections to reduce number of calculations possible_connections = np.unique(all_connections, axis=1) possible_connections = np.sort(possible_connections, axis=0) possible_connections = np.unique(possible_connections, axis=1) possible_connections = possible_connections[:, possible_connections[0, :] != possible_connections[1, :]] if possible_connections.shape[1] < 1: return [], [] # precalculate the possible direction vectors possible_direction_vectors = vertices.data[possible_connections[0]] - vertices.data[possible_connections[1]] possible_direction_vectors = possible_direction_vectors / np.linalg.norm(possible_direction_vectors, axis=1)[:, np.newaxis] owned_lines_per_possible_connections = [list() for i in range(possible_connections.shape[1])] # assign lines to possible connections for line_idx, i, j in zip(line_idx_list, begin_vertex_list, end_vertex_list): if i == j: continue i, j = min(i, j), max(i, j) for connection_idx, connection in enumerate(possible_connections.T): if np.all((i, j) == connection): owned_lines_per_possible_connections[connection_idx].append(line_idx) break # check if the lines are in the same direction as the possible connection for fitted_line_idx, owned_lines_per_possible_connection in enumerate(owned_lines_per_possible_connections): line_deviations = np.abs( np.dot(line_directions[owned_lines_per_possible_connection], possible_direction_vectors[fitted_line_idx])) if np.any(line_deviations > deviation_threshold): connections.append(possible_connections[:, fitted_line_idx]) vertices = [{"xy": v, "type": "apex"} for v in apex_centroids] vertices += [{"xy": v, "type": "apex"} for v in missed_vertices] vertices += [{"xy": v, "type": "eave_end_point"} for v in eave_end_point_centroids] return vertices, connections def get_uv_depth(vertices, depth): '''Get the depth of the vertices from the depth image''' # depth[depth > 5000] = np.inf uv = np.array([v['xy'] for v in vertices]) uv_int = uv.astype(np.int32) H, W = depth.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 = depth[(uv_int[:, 1], uv_int[:, 0])] return uv, vertex_depth def merge_vertices_3d(vert_edge_per_image, merge_th=0.1, **kwargs): '''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 <= merge_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 = {} for count, idxs in enumerate(merged): new_vertices.append(all_3d_vertices[idxs].mean(axis=0)) for idx in idxs: old_idx_to_new[idx] = count # 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): '''Prune vertices that are not connected to any other vertex''' connected = defaultdict(list) for c in connections_3d: connected[c[0]].append(c) connected[c[1]].append(c) new_indexes = {} new_verts = [] connected_out = [] for k, v in connected.items(): vert = all_3d_vertices[k] if tuple(vert) not in new_verts: new_verts.append(tuple(vert)) new_indexes[k] = len(new_verts) - 1 for k, v in connected.items(): for vv in v: connected_out.append((new_indexes[vv[0]], new_indexes[vv[1]])) connected_out = list(set(connected_out)) return np.array(new_verts), connected_out def predict(entry, visualize=False, scale_estimation_coefficient=2.5, **kwargs) -> Tuple[np.ndarray, List[int]]: if 'gestalt' not in entry or 'depthcm' not in entry or 'K' not in entry or 'R' not in entry or 't' not in entry: print('Missing required fields in the entry') return (entry['__key__'], *empty_solution()) entry = hoho.decode(entry) vert_edge_per_image = {} image_dict = {} for k, v in entry["images"].items(): image_dict[v.name] = v points = [v.xyz for k, v in entry["points3d"].items()] too_big = len(points) > 25000 if not too_big: points = np.array(points) point_keys = [k for k, v in entry["points3d"].items()] point_keys = np.array(point_keys) # print(len(points)) clustered = shc.fclusterdata(points, 100, criterion='distance') clustered_indices = np.argsort(clustered) points = points[clustered_indices] point_keys = point_keys[clustered_indices] clustered = clustered[clustered_indices] _, cluster_indices = np.unique(clustered, return_index=True) clustered_points = np.split(points, cluster_indices[1:]) clustered_keys = np.split(point_keys, cluster_indices[1:]) biggest_cluster_index = np.argmax([len(i) for i in clustered_points]) # biggest_cluster = clustered_points[biggest_cluster_index] biggest_cluster_keys = clustered_keys[biggest_cluster_index] biggest_cluster_keys = set(biggest_cluster_keys) for i, (gest, depthcm, K, R, t, imagekey) in enumerate(zip(entry['gestalt'], entry['depthcm'], entry['K'], entry['R'], entry['t'], entry['__imagekey__'] )): try: gest_seg = gest.resize(depthcm.size) gest_seg_np = np.array(gest_seg).astype(np.uint8) vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, **kwargs) 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 belonging_points = [] for i in image_dict[imagekey].point3D_ids[np.where(image_dict[imagekey].point3D_ids != -1)]: if not too_big: if i in biggest_cluster_keys: belonging_points.append(entry["points3d"][i]) else: belonging_points.append(entry["points3d"][i]) if len(belonging_points) < 1: print(f'No 3D points in image {i}') vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3)) raise KeyError projected2d, _ = cv2.projectPoints(np.array([i.xyz for i in belonging_points]), R, t, K, 0) important = np.where(np.all(projected2d >= 0, axis=2)) # Normalize the uv to the camera intrinsics world_to_cam = np.eye(4) world_to_cam[:3, :3] = R world_to_cam[:3, 3] = t homo_belonging_points = cv2.convertPointsToHomogeneous(np.array([i.xyz for i in belonging_points])) depth = cv2.convertPointsFromHomogeneous(cv2.transform(homo_belonging_points, world_to_cam)) depth = np.array([i[0][2] for i in depth]) depth = depth[important[0]] projected2d = projected2d[important] if len(depth) < 1: print(f'No 3D points in image {i}') vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3)) raise KeyError # print(projected2d.shape, depth.shape) interpolator = si.NearestNDInterpolator(projected2d, depth) vertex_coordinates = np.array([v['xy'] for v in vertices]) xi, yi = vertex_coordinates[:, 0], vertex_coordinates[:, 1] depth_vert = interpolator(xi, yi) xy_local = np.ones((len(vertex_coordinates), 3)) xy_local[:, 0] = (vertex_coordinates[:, 0] - K[0, 2]) / K[0, 0] xy_local[:, 1] = (vertex_coordinates[:, 1] - K[1, 2]) / K[1, 1] # Get the 3D vertices vertices_3d_local = depth_vert[..., None] * (xy_local / np.linalg.norm(xy_local, axis=1)[..., None]) 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) except KeyError: gest_seg = gest.resize(depthcm.size) gest_seg_np = np.array(gest_seg).astype(np.uint8) # Metric3D depth_np = np.array(depthcm) / scale_estimation_coefficient vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, **kwargs) 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_np) # Normalize the uv to the camera intrinsics xy_local = np.ones((len(uv), 3)) 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 = depth_vert[..., None] * (xy_local / np.linalg.norm(xy_local, axis=1)[..., None]) 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, **kwargs) all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d # all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d) 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 (entry['__key__'], *empty_solution()) if visualize: from hoho.viz3d import plot_estimate_and_gt plot_estimate_and_gt(all_3d_vertices_clean, connections_3d_clean, entry['wf_vertices'], entry['wf_edges']) return entry['__key__'], all_3d_vertices_clean, connections_3d_clean