import os import json import cv2 import numpy as np from collections import defaultdict from scipy import ndimage def generate_graph(annot, image_path, out_path): lines = annot['lines'] junctions = annot['junctions'] line_junc_mat = np.array(annot['lineJunctionMatrix']) planes = annot['planes'] plane_line_mat = annot['planeLineMatrix'] plane_to_line = np.array(annot['planeLineMatrix']) line_to_plane = plane_to_line.T semantics = annot['semantics'] all_room_edges = get_room_edges(semantics, planes, lines, junctions, plane_to_line, line_junc_mat) all_room_edges = filter_rooms(all_room_edges, im_size=256) all_colinear_pairs = find_all_colinear_paris(all_room_edges) colinear_sets = combine_colinear_edges(all_colinear_pairs) for colinear_set in colinear_sets: edges_to_merge = list(colinear_set) edges_to_merge = sorted(edges_to_merge, key=lambda x: -x[0]) merged_edges = merge_edges(edges_to_merge) for merged_edge, old_edge in zip(merged_edges, edges_to_merge): if len(merged_edge) > 0: assert merged_edge[0][0] == old_edge[0] room_idx = merged_edge[0][0] if len(merged_edge) == 1 and merged_edge[0] == old_edge: continue # change room graph accordingly replaced_idx = all_room_edges[room_idx].index(old_edge[1]) all_room_edges[room_idx].pop(replaced_idx) for new_idx, new_edge in enumerate(merged_edge): insert_idx = new_idx + replaced_idx all_room_edges[room_idx].insert(insert_idx, new_edge[1]) else: room_idx = old_edge[0] replaced_idx = all_room_edges[room_idx].index(old_edge[1]) all_room_edges[room_idx].pop(replaced_idx) # take intersection for every rooms to recover the room structure refined_room_edges = [adjust_room_edges(room_edges) for room_edges in all_room_edges] # clean every room loop by removing I-shape corners cleaned_room_edges = clean_room_edges(refined_room_edges) global_graph = defaultdict(list) for room_edges in cleaned_room_edges: for edge in room_edges: c1, c2 = edge global_graph[c1] += [c2, ] global_graph[c2] += [c1, ] for corner in global_graph: global_graph[corner] = list(set(global_graph[corner])) annot_path = os.path.join(out_path, 'annot.npy') np.save(annot_path, global_graph) # draw the planar graph on the density map viz_image = cv2.imread(image_path) for c, connections in global_graph.items(): for other_c in connections: cv2.line(viz_image, (int(c[0]), int(c[1])), (int(other_c[0]), int(other_c[1])), (255, 255, 0), 2) for c in global_graph.keys(): cv2.circle(viz_image, (int(c[0]), int(c[1])), 3, (0, 0, 255), -1) # viz_image = np.zeros([256, 266, 3]).astype(np.uint8) # room_idx = 0 # for room_edges in cleaned_room_edges: # for line_idx, edge in enumerate(room_edges): # c1, c2 = np.array(edge).astype(np.int) # cv2.line(viz_image, tuple(c1), tuple(c2), (255, 255, 0), 2) # cv2.circle(viz_image, tuple(c1), 3, (0, 0, 255), -1) # cv2.circle(viz_image, tuple(c2), 3, (0, 0, 255), -1) # room_idx += 1 cv2.imwrite(os.path.join(out_path, 'planar_graph.png'), viz_image) def get_room_edges(semantics, planes, lines, junctions, plane_to_line, line_junc_mat): room_edges = list() for semantic in semantics: plane_ids = semantic['planeID'] label = semantic['type'] if label in ['door', 'window', 'outwall']: # skip non-room elements continue all_planes = [planes[idx] for idx in plane_ids] floor_planes = [plane for plane in all_planes if plane['type'] == 'floor'] assert len(floor_planes) == 1, 'There should be only one floor for each room' floor_plane = floor_planes[0] floor_plane_id = floor_plane['ID'] line_ids = np.where(plane_to_line[floor_plane_id])[0].tolist() floor_lines_raw = [lines[line_id] for line_id in line_ids] floor_lines = list() for line_idx, floor_line in enumerate(floor_lines_raw): c_id_1, c_id_2 = np.where(line_junc_mat[floor_line['ID']])[0].tolist() c1 = tuple(junctions[c_id_1]['coordinate'][:2]) c2 = tuple(np.array(junctions[c_id_2]['coordinate'][:2])) if c1 == c2: continue floor_lines.append((c1, c2)) floor_lines = list(set(floor_lines)) # remove duplications floor_lines = sort_room_edges(floor_lines) room_edges.append(floor_lines) return room_edges def sort_room_edges(lines): cur_id = 0 picked = [False] * len(lines) id_list = [0, ] while len(id_list) < len(lines): line = lines[cur_id] picked[cur_id] = True check_ = [(line[1] in other) and not picked[other_idx] for other_idx, other in enumerate(lines)] next_ids = np.nonzero(check_)[0] try: assert len(next_ids) == 1 except: raise WrongRoomError('Invalid room shape') next_id = next_ids[0] id_list.append(next_id) if lines[next_id][1] == line[1]: # swap the two endpoints, to make the loop valid. lines[next_id] = (lines[next_id][1], lines[next_id][0]) cur_id = next_id if lines[next_id][1] == lines[0][0]: # already form a closed loop, then skip the remaining lines break sorted_lines = [lines[idx] for idx in id_list] return sorted_lines def find_all_colinear_paris(all_room_edges): colinear_pairs = list() for room_idx, room_edges in enumerate(all_room_edges): for edge_idx, edge in enumerate(room_edges): for other_room_idx, other_edges in enumerate(all_room_edges): if other_room_idx < room_idx: continue for other_edge_idx, other_edge in enumerate(other_edges): if other_room_idx == room_idx and other_edge_idx <= edge_idx: continue if _check_colinear(edge, other_edge, line_dist_th=8): ele1 = (room_idx, edge) ele2 = (other_room_idx, other_edge) colinear_pairs.append([ele1, ele2]) return colinear_pairs def combine_colinear_edges(colinear_pairs): all_colinear_sets = list() all_pairs = list(colinear_pairs) # make a copy of the input list combined = [False] * len(colinear_pairs) while len(all_pairs) > 0: colinear_set = _combine_colinear_pairs(0, all_pairs, combined) all_colinear_sets.append(colinear_set) all_pairs = [all_pairs[i] for i in range(len(all_pairs)) if combined[i] is False] combined = [False] * len(all_pairs) return all_colinear_sets def _combine_colinear_pairs(idx, all_pairs, combined): colinear_set = set(all_pairs[idx]) combined[idx] = True for other_idx, pair in enumerate(all_pairs): if not combined[other_idx] and ( all_pairs[idx][0] in all_pairs[other_idx] or all_pairs[idx][1] in all_pairs[other_idx]): colinear_set = colinear_set.union(_combine_colinear_pairs(other_idx, all_pairs, combined)) return colinear_set def _check_colinear(e1, e2, line_dist_th=8): # first check whether two line segments are parallel to each other, if not, return False directly len_e1 = len_edge(e1) len_e2 = len_edge(e2) # we need to always make e2 the shorter one if len_e1 < len_e2: e1, e2 = e2, e1 v1_01 = (e1[1][0] - e1[0][0], e1[1][1] - e1[0][1]) v1_10 = (e1[0][0] - e1[1][0], e1[0][1] - e1[1][1]) v2_01 = (e2[1][0] - e2[0][0], e2[1][1] - e2[0][1]) v2_10 = (e2[0][0] - e2[1][0], e2[0][1] - e2[1][1]) len_1 = np.sqrt(v1_01[0] ** 2 + v1_01[1] ** 2) len_2 = np.sqrt(v2_01[0] ** 2 + v2_01[1] ** 2) if len_1 == 0 or len_2 == 0: cos = 0 else: cos = (v1_01[0] * v2_01[0] + v1_01[1] * v2_01[1]) / (len_1 * len_2) if abs(cos) > 0.99: # then check the distance between two parallel lines len_10_20 = len_edge((e1[0], e2[0])) len_10_21 = len_edge((e1[0], e2[1])) len_11_20 = len_edge((e1[1], e2[0])) len_11_21 = len_edge((e1[1], e2[1])) # two endpoints are very close, then we can say these two edges are colinear if np.min([len_10_20, len_10_21, len_11_20, len_11_21]) <= 5: return True # otherwise we need to check the distance first v_10_20 = (e2[0][0] - e1[0][0], e2[0][1] - e1[0][1]) cos_11_10_20 = (v1_01[0] * v_10_20[0] + v1_01[1] * v_10_20[1]) / (len_1 * len_10_20) sin_11_10_20 = np.sqrt(1 - cos_11_10_20 ** 2) dist_20_e1 = len_10_20 * sin_11_10_20 if dist_20_e1 <= line_dist_th: # we need two check whether they have some overlaps v_11_20 = (e2[0][0] - e1[1][0], e2[0][1] - e1[1][1]) cos_10_11_20 = (v1_10[0] * v_11_20[0] + v1_10[1] * v_11_20[1]) / (len_1 * len_11_20) if cos_11_10_20 >= 0 and cos_10_11_20 >= 0: return True v_10_21 = (e2[1][0] - e1[0][0], e2[1][1] - e1[0][1]) cos_11_10_21 = (v1_01[0] * v_10_21[0] + v1_01[1] * v_10_21[1]) / (len_1 * len_10_21) v_11_21 = (e2[1][0] - e1[1][0], e2[1][1] - e1[1][1]) cos_10_11_21 = (v1_10[0] * v_11_21[0] + v1_10[1] * v_11_21[1]) / (len_1 * len_11_21) if cos_11_10_21 >= 0 and cos_10_11_21 >= 0: return True return False else: # if the two line segments have distance > 3, we can say they are not colinear return False else: return False def merge_edges(edges): base_e = edges[0][1] merged_edges = [edges[0], ] base_len = np.sqrt((base_e[1][0] - base_e[0][0]) ** 2 + (base_e[1][1] - base_e[0][1]) ** 2) base_unit_v = ((base_e[1][0] - base_e[0][0]) / base_len, (base_e[1][1] - base_e[0][1]) / base_len) for edge in edges[1:]: room_idx = edge[0] e = edge[1] v_b0e0 = (e[0][0] - base_e[0][0], e[0][1] - base_e[0][1]) proj_len = (v_b0e0[0] * base_unit_v[0] + v_b0e0[1] * base_unit_v[1]) proj_e0 = (int(base_e[0][0] + base_unit_v[0] * proj_len), int(base_e[0][1] + base_unit_v[1] * proj_len)) proj_e1 = (int(proj_e0[0] + e[1][0] - e[0][0]), int(proj_e0[1] + e[1][1] - e[0][1])) new_e = (proj_e0, proj_e1) new_edge = (room_idx, new_e) merged_edges.append(new_edge) adjusted_merged_edges = adjust_colinear_edges(merged_edges) return adjusted_merged_edges def adjust_colinear_edges(edges): base_corner = (edges[0][0], edges[0][1][0]) all_corners = [base_corner, (edges[0][0], edges[0][1][1])] for edge in edges[1:]: all_corners.append((edge[0], edge[1][0])) all_corners.append((edge[0], edge[1][1])) unit_v = unit_v_edge(edges[0][1]) corner_projs = list() # FIXME: need to fix the corner coords here, it's wrong now!! they are not merged... for room, other_c in all_corners: v_base_c = (other_c[0] - base_corner[1][0], other_c[1] - base_corner[1][1]) proj = (unit_v[0] * v_base_c[0] + unit_v[1] * v_base_c[1]) corner_projs.append(proj) order = np.argsort(corner_projs).tolist() # # merge corners that are too close to the prev corner # for o_idx, corner_idx in enumerate(order[1:]): # corner = all_corners[corner_idx][1] # prev_idx = order[o_idx] # prev_corner = all_corners[prev_idx][1] # dist = len_edge((corner, prev_corner)) # if dist <= 5: # all_corners[corner_idx] = (all_corners[corner_idx][0], prev_corner) adjusted_edges = list() for idx, edge in enumerate(edges): room_idx = edge[0] idx_1 = idx * 2 idx_2 = idx * 2 + 1 adj_idx_1 = order.index(idx_1) adj_idx_2 = order.index(idx_2) step_direction = 1 if adj_idx_2 > adj_idx_1 else -1 adjusted_edge = list() for o_idx in range(adj_idx_1, adj_idx_2, step_direction): c_idx = order[o_idx] next_c_idx = order[o_idx + step_direction] segment = (room_idx, (all_corners[c_idx][1], all_corners[next_c_idx][1])) if len_edge(segment[1]) == 0: continue adjusted_edge.append(segment) adjusted_edges.append(adjusted_edge) return adjusted_edges def adjust_room_edges(room_edges): refined_room_edges = list() init_room_edges = list(room_edges) for edge_i, edge in enumerate(room_edges): next_i = edge_i while True: next_i = next_i + 1 if next_i < len(room_edges) - 1 else 0 next_edge = room_edges[next_i] if next_edge[0] != next_edge[1]: break if edge[1] == next_edge[0]: # no need for refining refined_room_edges.append(edge) else: # the two corners disagree, refinement is required if edge[0] == edge[1]: print('skip collasped edge') continue unit_edge = unit_v_edge(edge) ext_edge = ((edge[0][0] - unit_edge[0] * 50, edge[0][1] - unit_edge[1] * 50), (edge[1][0] + unit_edge[0] * 50, edge[1][1] + unit_edge[1] * 50)) unit_next = unit_v_edge(next_edge) ext_next = ((next_edge[0][0] - unit_next[0] * 50, next_edge[0][1] - unit_next[1] * 50), (next_edge[1][0] + unit_next[0] * 50, next_edge[1][1] + unit_next[1] * 50)) intersec = get_intersection(ext_edge[0], ext_edge[1], ext_next[0], ext_next[1]) try: assert intersec is not None except: print('no intersect, move endpoint directly') intersec = next_edge[0] intersec = (int(np.round(intersec[0])), int(np.round(intersec[1]))) refined_edge = (edge[0], intersec) refined_room_edges.append(refined_edge) room_edges[edge_i] = refined_edge room_edges[next_i] = (intersec, next_edge[1]) if next_i < edge_i: refined_room_edges[next_i] = room_edges[next_i] # drop collapsed edges refined_room_edges = [edge for edge in refined_room_edges if edge[0] != edge[1]] for edge_i in range(len(refined_room_edges)): next_i = edge_i + 1 if edge_i < len(refined_room_edges) - 1 else 0 if refined_room_edges[edge_i][1] != refined_room_edges[next_i][0]: new_edge = (refined_room_edges[edge_i][0], refined_room_edges[next_i][0]) refined_room_edges[edge_i] = new_edge return refined_room_edges def clean_room_edges(all_room_edges): refined_room_paths = [_extract_room_path(room_edges) for room_edges in all_room_edges] corner_to_room = defaultdict(list) for room_idx, room_path in enumerate(refined_room_paths): for corner in room_path: corner_to_room[corner].append(room_idx) # remove I-shape corner used by only one room for room_idx, room_edges in enumerate(all_room_edges): cp_room_edges = list(room_edges) rm_flag = True while rm_flag: rm_flag = False for edge_i, edge in enumerate(cp_room_edges): prev_i = edge_i - 1 prev_edge = cp_room_edges[prev_i] if _check_colinear(prev_edge, edge, line_dist_th=5): rm_candidate = edge[0] if len(corner_to_room[rm_candidate]) == 1 and corner_to_room[rm_candidate][0] == room_idx: cp_room_edges[prev_i] = (prev_edge[0], edge[1]) rm_flag = True cp_room_edges.pop(edge_i) break next_i = edge_i + 1 if edge_i < len(cp_room_edges) - 1 else 0 next_edge = cp_room_edges[next_i] if _check_colinear(next_edge, edge, line_dist_th=5): rm_candidate = edge[1] if len(corner_to_room[rm_candidate]) == 1 and corner_to_room[rm_candidate][0] == room_idx: cp_room_edges[next_i] = (edge[0], next_edge[1]) rm_flag = True cp_room_edges.pop(edge_i) break if len(cp_room_edges) != len(room_edges): all_room_edges[room_idx] = cp_room_edges corner_to_room = get_corner_to_room(all_room_edges) all_corners = list(corner_to_room.keys()) corners_to_merge = find_corners_to_merge(all_corners, corner_to_room) while corners_to_merge is not None: num_aff = [len(corner_to_room[x]) for x in corners_to_merge] order = np.argsort(num_aff)[::-1] base_corner = corners_to_merge[order[0]] for corner in corners_to_merge: if corner == base_corner: continue all_room_edges = move_corner(corner, base_corner, corner_to_room, all_room_edges) corner_to_room = get_corner_to_room(all_room_edges) all_corners = list(corner_to_room.keys()) corners_to_merge = find_corners_to_merge(all_corners, corner_to_room) # for room_idx, room_edges in enumerate(all_room_edges): # cp_room_edges = list(room_edges) # rm_flag = True # while rm_flag: # rm_flag = False # for edge_i, edge in enumerate(cp_room_edges): # len_e = len_edge(edge) # if len_e <= 5: # if len(corner_to_room[edge[0]]) == 1: # prev_i = edge_i - 1 # prev_edge = cp_room_edges[prev_i] # cp_room_edges[prev_i] = (prev_edge[0], edge[1]) # rm_flag = True # cp_room_edges.pop(edge_i) # break # elif len(corner_to_room[edge[1]]) == 1: # next_i = edge_i + 1 if edge_i < len(cp_room_edges) - 1 else 0 # next_edge = cp_room_edges[next_i] # cp_room_edges[next_i] = (edge[0], next_edge[1]) # rm_flag = True # cp_room_edges.pop(edge_i) # else: # continue # # if len(cp_room_edges) != len(room_edges): # all_room_edges[room_idx] = cp_room_edges return all_room_edges def move_corner(c, target, corner_to_room, all_room_edges): rooms = corner_to_room[c] for room_idx in rooms: for edge_idx, edge in enumerate(all_room_edges[room_idx]): if c in edge: if c == edge[0]: new_edge = (target, edge[1]) elif c == edge[1]: new_edge = (edge[0], target) else: continue all_room_edges[room_idx][edge_idx] = new_edge return all_room_edges def find_corners_to_merge(all_corners, corner_to_room, th=5): all_close_pairs = list() for idx1, corner in enumerate(all_corners): for idx2, other_corner in enumerate(all_corners): if idx2 <= idx1: continue if len_edge((corner, other_corner)) <= th: rooms_1 = tuple(sorted(corner_to_room[corner])) rooms_2 = tuple(sorted(corner_to_room[other_corner])) if rooms_1 == rooms_2: continue elif len(rooms_1) ==1: if rooms_1[0] in list(rooms_2): continue else: all_close_pairs.append([corner, other_corner]) elif len(rooms_2) ==1: if rooms_2[0] in list(rooms_1): continue else: all_close_pairs.append([corner, other_corner]) else: all_close_pairs.append([corner, other_corner]) if len(all_close_pairs) == 0: return None close_set = find_one_close_set(all_close_pairs) corners_to_merge = list(close_set) return corners_to_merge def find_one_close_set(all_corner_paris): all_pairs = list(all_corner_paris) # make a copy of the input list combined = [False] * len(all_corner_paris) close_set = _combine_colinear_pairs(0, all_pairs, combined) return close_set def get_corner_to_room(all_room_edges): room_paths = [_extract_room_path(room_edges) for room_edges in all_room_edges] corner_to_room = defaultdict(list) for room_idx, room_path in enumerate(room_paths): for corner in room_path: corner_to_room[corner].append(room_idx) return corner_to_room def filter_rooms(all_room_edges, im_size): # filter rooms that are covered by larger rooms room_masks = list() updated_room_edges = list() for room_edges in all_room_edges: room_mask = draw_room_seg_from_edges(room_edges, im_size) if room_mask is not None and room_mask.sum() > 20: # remove too small rooms room_masks.append(room_mask) updated_room_edges.append(room_edges) all_room_edges = updated_room_edges removed = list() for room_idx, room_mask in enumerate(room_masks): # do not consider the current room, and do not consider removed rooms other_masks = [room_masks[i] for i in range(len(all_room_edges)) if i != room_idx and i not in removed] if len(other_masks) == 0: # if all other masks are removed.. other_masks_all = np.zeros([im_size, im_size]) else: other_masks_all = np.clip(np.sum(np.stack(other_masks, axis=-1), axis=-1), 0, 1) joint_mask = np.clip(other_masks_all + room_mask, 0, 1) mask_area = room_mask.sum() overlap_area = mask_area + other_masks_all.sum() - joint_mask.sum() if overlap_area / mask_area > 0.5: removed.append(room_idx) all_room_edges = [all_room_edges[idx] for idx in range(len(all_room_edges)) if idx not in removed] return all_room_edges ## Utils class WrongRoomError(Exception): pass def _extract_room_path(room_edges): room_path = [edge[0] for edge in room_edges] return room_path def len_edge(e): return np.sqrt((e[1][0] - e[0][0]) ** 2 + (e[1][1] - e[0][1]) ** 2) def unit_v_edge(e): len_e = len_edge(e) assert len_e != 0 unit_v = ((e[1][0] - e[0][0]) / len_e, (e[1][1] - e[0][1]) / len_e) return unit_v def get_intersection(p0, p1, p2, p3): """ reference: StackOverflow https://stackoverflow.com/questions/563198/how-do-you-detect-where-two-line-segments-intersect#565282 """ s1_x = p1[0] - p0[0] s1_y = p1[1] - p0[1] s2_x = p3[0] - p2[0] s2_y = p3[1] - p2[1] s = (-s1_y * (p0[0] - p2[0]) + s1_x * (p0[1] - p2[1])) / (-s2_x * s1_y + s1_x * s2_y) t = (s2_x * (p0[1] - p2[1]) - s2_y * (p0[0] - p2[0])) / (-s2_x * s1_y + s1_x * s2_y) if 1 >= s >= 0 and 1 >= t >= 0: i_x = p0[0] + (t * s1_x) i_y = p0[1] + (t * s1_y) return (i_x, i_y) else: return None def draw_room_seg_from_edges(edges, im_size): edge_map = np.zeros([im_size, im_size]) for edge in edges: edge = np.array(edge).astype(np.int) cv2.line(edge_map, tuple(edge[0]), tuple(edge[1]), 1, 3) reverse_edge_map = 1 - edge_map label, num_features = ndimage.label(reverse_edge_map) if num_features < 2: return None bg_label = label[0, 0] num_labels = [(label==l).sum() for l in range(1, num_features+1)] num_labels[bg_label-1] = 0 room_label = np.argmax(num_labels) + 1 room_map = np.zeros([im_size, im_size]) room_map[np.where(label == room_label)] = 1 return room_map if __name__ == '__main__': data_base = './montefloor_data/' dir_names = list(sorted(os.listdir(data_base))) invalid_scenes = list() for dir_name in dir_names: if 'scene' not in dir_name: continue data_dir = os.path.join(data_base, dir_name) annot_path = os.path.join(data_dir, 'annotation_3d.json') with open(annot_path) as f: annot = json.load(f) image_path = os.path.join(data_dir, 'density.png') try: generate_graph(annot, image_path, data_dir) except WrongRoomError: invalid_scenes.append(dir_name) print('Finish processing data {}'.format(dir_name)) print('Failed on {} scenes with invalid rooms: {}'.format(len(invalid_scenes), invalid_scenes))