| """ |
| This code is an adaptation that uses Structured 3D for the code base. |
| |
| Reference: https://github.com/bertjiazheng/Structured3D |
| """ |
|
|
| import numpy as np |
| import cv2 |
| from shapely.geometry import Polygon |
| import random |
|
|
| def parse_floor_plan_polys(annos): |
| planes = [] |
| for semantic in annos['semantics']: |
| for planeID in semantic['planeID']: |
| if annos['planes'][planeID]['type'] == 'floor': |
| planes.append({'planeID': planeID, 'type': semantic['type']}) |
|
|
| if semantic['type'] == 'outwall': |
| outerwall_planes = semantic['planeID'] |
|
|
| |
| lines_holes = [] |
| for semantic in annos['semantics']: |
| if semantic['type'] in ['window', 'door']: |
| for planeID in semantic['planeID']: |
| lines_holes.extend(np.where(np.array(annos['planeLineMatrix'][planeID]))[0].tolist()) |
| lines_holes = np.unique(lines_holes) |
|
|
| |
| junctions = np.array([junc['coordinate'] for junc in annos['junctions']]) |
| junction_floor = np.where(np.isclose(junctions[:, -1], 0))[0] |
|
|
| |
| polygons = [] |
| for plane in planes: |
| lineIDs = np.where(np.array(annos['planeLineMatrix'][plane['planeID']]))[0].tolist() |
| junction_pairs = [np.where(np.array(annos['lineJunctionMatrix'][lineID]))[0].tolist() for lineID in lineIDs] |
| polygon = convert_lines_to_vertices(junction_pairs) |
| polygons.append([polygon[0], plane['type']]) |
|
|
| outerwall_floor = [] |
| for planeID in outerwall_planes: |
| lineIDs = np.where(np.array(annos['planeLineMatrix'][planeID]))[0].tolist() |
| lineIDs = np.setdiff1d(lineIDs, lines_holes) |
| junction_pairs = [np.where(np.array(annos['lineJunctionMatrix'][lineID]))[0].tolist() for lineID in lineIDs] |
| for start, end in junction_pairs: |
| if start in junction_floor and end in junction_floor: |
| outerwall_floor.append([start, end]) |
|
|
| outerwall_polygon = convert_lines_to_vertices(outerwall_floor) |
| polygons.append([outerwall_polygon[0], 'outwall']) |
|
|
| return polygons |
|
|
| def convert_lines_to_vertices(lines): |
| """ |
| convert line representation to polygon vertices |
| |
| """ |
| polygons = [] |
| lines = np.array(lines) |
|
|
| polygon = None |
| while len(lines) != 0: |
| if polygon is None: |
| polygon = lines[0].tolist() |
| lines = np.delete(lines, 0, 0) |
|
|
| lineID, juncID = np.where(lines == polygon[-1]) |
| vertex = lines[lineID[0], 1 - juncID[0]] |
| lines = np.delete(lines, lineID, 0) |
|
|
| if vertex in polygon: |
| polygons.append(polygon) |
| polygon = None |
| else: |
| polygon.append(vertex) |
|
|
| return polygons |
|
|
| def generate_floorplan(annos, polygons, height, width, ignore_types, include_types=None, fillpoly=True, constant_color=False, shuffle=False): |
| """ |
| plot floorplan |
| |
| """ |
|
|
| floor_map = np.zeros((height, width)) |
|
|
| junctions = np.array([junc['coordinate'][:2] for junc in annos['junctions']]) |
|
|
| room_ind = 0 |
| if shuffle: |
| room_ind = np.random.randint(0, 2) |
|
|
| polygons_list = [] |
| for poly_ind, (polygon, poly_type) in enumerate(polygons): |
| if poly_type in ignore_types: |
| continue |
| if include_types is not None and poly_type not in include_types: |
| continue |
|
|
| polygon = junctions[np.array(polygon)].astype(np.int32) |
|
|
| poly_shapely = Polygon(polygon) |
| area = poly_shapely.area |
|
|
| |
| if area < 100: |
| continue |
|
|
| polygons_list.append(polygon) |
|
|
| if shuffle: |
| random.shuffle(polygons_list) |
| for poly_ind, polygon in enumerate(polygons_list): |
|
|
| if shuffle: |
| room_ind += np.random.randint(1, 2) |
| else: |
| room_ind += 1 |
|
|
| if fillpoly: |
| if constant_color: |
| clr = 1. |
| else: |
| clr = room_ind |
| cv2.fillPoly(floor_map, [polygon], color=clr) |
| else: |
| assert constant_color |
| cv2.polylines(floor_map, [polygon], isClosed=True, color=1., thickness=3) |
|
|
| return floor_map, polygons_list |
