import os import json import argparse import open3d import pymesh import numpy as np import matplotlib.pyplot as plt from shapely.geometry import Polygon from descartes.patch import PolygonPatch from misc.figures import plot_coords from misc.colors import colormap_255, semantics_cmap def visualize_wireframe(annos): """visualize wireframe """ colormap = np.array(colormap_255) / 255 junctions = np.array([item['coordinate'] for item in annos['junctions']]) _, junction_pairs = np.where(np.array(annos['lineJunctionMatrix'])) junction_pairs = junction_pairs.reshape(-1, 2) # extract hole lines 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) # extract cuboid lines cuboid_lines = [] for cuboid in annos['cuboids']: for planeID in cuboid['planeID']: cuboid_lineID = np.where(np.array(annos['planeLineMatrix'][planeID]))[0].tolist() cuboid_lines.extend(cuboid_lineID) cuboid_lines = np.unique(cuboid_lines) cuboid_lines = np.setdiff1d(cuboid_lines, lines_holes) # visualize junctions connected_junctions = junctions[np.unique(junction_pairs)] connected_colors = np.repeat(colormap[0].reshape(1, 3), len(connected_junctions), axis=0) junction_set = open3d.geometry.PointCloud() junction_set.points = open3d.utility.Vector3dVector(connected_junctions) junction_set.colors = open3d.utility.Vector3dVector(connected_colors) # visualize line segments line_colors = np.repeat(colormap[5].reshape(1, 3), len(junction_pairs), axis=0) # color holes if len(lines_holes) != 0: line_colors[lines_holes] = colormap[6] # color cuboids if len(cuboid_lines) != 0: line_colors[cuboid_lines] = colormap[2] line_set = open3d.geometry.LineSet() line_set.points = open3d.utility.Vector3dVector(junctions) line_set.lines = open3d.utility.Vector2iVector(junction_pairs) line_set.colors = open3d.utility.Vector3dVector(line_colors) open3d.visualization.draw_geometries([junction_set, line_set]) def project(x, meta): """ project 3D to 2D for polygon clipping """ proj_axis = max(range(3), key=lambda i: abs(meta['normal'][i])) return tuple(c for i, c in enumerate(x) if i != proj_axis) def project_inv(x, meta): """ recover 3D points from 2D """ # Returns the vector w in the walls' plane such that project(w) equals x. proj_axis = max(range(3), key=lambda i: abs(meta['normal'][i])) w = list(x) w[proj_axis:proj_axis] = [0.0] c = -meta['offset'] for i in range(3): c -= w[i] * meta['normal'][i] c /= meta['normal'][proj_axis] w[proj_axis] = c return tuple(w) def triangulate(points): """ triangulate the plane for operation and visualization """ num_points = len(points) indices = np.arange(num_points, dtype=np.int) segments = np.vstack((indices, np.roll(indices, -1))).T tri = pymesh.triangle() tri.points = np.array(points) tri.segments = segments tri.verbosity = 0 tri.run() return tri.mesh def clip_polygon(polygons, vertices_hole, junctions, meta): """ clip polygon the hole """ if len(polygons) == 1: junctions = [junctions[vertex] for vertex in polygons[0]] mesh_wall = triangulate(junctions) vertices = np.array(mesh_wall.vertices) faces = np.array(mesh_wall.faces) return vertices, faces else: wall = [] holes = [] for polygon in polygons: if np.any(np.intersect1d(polygon, vertices_hole)): holes.append(polygon) else: wall.append(polygon) # extract junctions on this plane indices = [] junctions_wall = [] for plane in wall: for vertex in plane: indices.append(vertex) junctions_wall.append(junctions[vertex]) junctions_holes = [] for plane in holes: junctions_hole = [] for vertex in plane: indices.append(vertex) junctions_hole.append(junctions[vertex]) junctions_holes.append(junctions_hole) junctions_wall = [project(x, meta) for x in junctions_wall] junctions_holes = [[project(x, meta) for x in junctions_hole] for junctions_hole in junctions_holes] mesh_wall = triangulate(junctions_wall) for hole in junctions_holes: mesh_hole = triangulate(hole) mesh_wall = pymesh.boolean(mesh_wall, mesh_hole, 'difference') vertices = [project_inv(vertex, meta) for vertex in mesh_wall.vertices] return vertices, np.array(mesh_wall.faces) def draw_geometries_with_back_face(geometries): vis = open3d.visualization.Visualizer() vis.create_window() render_option = vis.get_render_option() render_option.mesh_show_back_face = True for geometry in geometries: vis.add_geometry(geometry) vis.run() vis.destroy_window() 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 visualize_plane(annos, args, eps=0.9): """visualize plane """ colormap = np.array(colormap_255) / 255 junctions = [item['coordinate'] for item in annos['junctions']] if args.color == 'manhattan': manhattan = dict() for planes in annos['manhattan']: for planeID in planes['planeID']: manhattan[planeID] = planes['ID'] # extract hole vertices 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) _, vertices_holes = np.where(np.array(annos['lineJunctionMatrix'])[lines_holes]) vertices_holes = np.unique(vertices_holes) # load polygons polygons = [] for semantic in annos['semantics']: for planeID in semantic['planeID']: plane_anno = annos['planes'][planeID] lineIDs = np.where(np.array(annos['planeLineMatrix'][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) vertices, faces = clip_polygon(polygon, vertices_holes, junctions, plane_anno) polygons.append([vertices, faces, planeID, plane_anno['normal'], plane_anno['type'], semantic['type']]) plane_set = [] plane_set_types = [] for i, (vertices, faces, planeID, normal, plane_type, semantic_type) in enumerate(polygons): # ignore the room ceiling # if semantic_type not in ['door', 'window']: # if plane_type == 'ceiling' and semantic_type not in ['door', 'window']: if semantic_type in ['door', 'window']: continue plane_vis = open3d.geometry.TriangleMesh() plane_vis.vertices = open3d.utility.Vector3dVector(vertices) plane_vis.triangles = open3d.utility.Vector3iVector(faces) if args.color == 'normal': if np.dot(normal, [1, 0, 0]) > eps: plane_vis.paint_uniform_color(colormap[0]) elif np.dot(normal, [-1, 0, 0]) > eps: plane_vis.paint_uniform_color(colormap[1]) elif np.dot(normal, [0, 1, 0]) > eps: plane_vis.paint_uniform_color(colormap[2]) elif np.dot(normal, [0, -1, 0]) > eps: plane_vis.paint_uniform_color(colormap[3]) elif np.dot(normal, [0, 0, 1]) > eps: plane_vis.paint_uniform_color(colormap[4]) elif np.dot(normal, [0, 0, -1]) > eps: plane_vis.paint_uniform_color(colormap[5]) else: plane_vis.paint_uniform_color(colormap[6]) elif args.color == 'manhattan': # paint each plane with manhattan world if planeID not in manhattan.keys(): plane_vis.paint_uniform_color(colormap[6]) else: plane_vis.paint_uniform_color(colormap[manhattan[planeID]]) plane_set.append(plane_vis) plane_set_types.append(plane_type) draw_geometries_with_back_face(plane_set) save_path = args.savepath if save_path: import sem_seg_utils as ss triangle_pcd_np_list = [] triangle_pcd_colors_np_list = [] for plane_triangle, plane_type in zip(plane_set, plane_set_types): triangle_pcd = plane_triangle.sample_points_uniformly(number_of_points=100) triangle_pcd_np = np.array(triangle_pcd.points) triangle_pcd_color_np = np.ones_like(triangle_pcd_np) * ss.class_name_to_id[plane_type] triangle_pcd_np_list.append(triangle_pcd_np) triangle_pcd_colors_np_list.append(triangle_pcd_color_np) final_pcd_np = np.concatenate(triangle_pcd_np_list, axis=0) final_pcd_colors_np = np.concatenate(triangle_pcd_colors_np_list, axis=0) o3d_final_pcd = open3d.geometry.PointCloud() o3d_final_pcd.points = open3d.utility.Vector3dVector(final_pcd_np) o3d_final_pcd.colors = open3d.utility.Vector3dVector(final_pcd_colors_np / 255.) scene_name = "scene_" + "{0:0=5}".format(args.scene) file_name = save_path + "/" + scene_name + "/" + scene_name + "_segmented.ply" open3d.io.write_point_cloud(file_name, o3d_final_pcd) def plot_floorplan(annos, polygons): """plot floorplan """ fig = plt.figure() ax = fig.add_subplot(1, 1, 1) junctions = np.array([junc['coordinate'][:2] for junc in annos['junctions']]) for (polygon, poly_type) in polygons: polygon = Polygon(junctions[np.array(polygon)]) plot_coords(ax, polygon.exterior, alpha=0.5) if poly_type == 'outwall': patch = PolygonPatch(polygon, facecolor=semantics_cmap[poly_type], alpha=0) else: patch = PolygonPatch(polygon, facecolor=semantics_cmap[poly_type], alpha=0.5) ax.add_patch(patch) plt.axis('equal') plt.axis('off') plt.show() def visualize_floorplan(annos): """visualize floorplan """ # extract the floor in each semantic for floorplan visualization 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'] # extract hole vertices 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 on the floor junctions = np.array([junc['coordinate'] for junc in annos['junctions']]) junction_floor = np.where(np.isclose(junctions[:, -1], 0))[0] # construct each polygon 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']) plot_floorplan(annos, polygons) def parse_args(): parser = argparse.ArgumentParser(description="Structured3D 3D Visualization") parser.add_argument("--path", required=True, help="dataset path", metavar="DIR") parser.add_argument("--scene", required=True, help="scene id", type=int) parser.add_argument("--type", choices=("floorplan", "wireframe", "plane"), default="plane", type=str) parser.add_argument("--color", choices=["normal", "manhattan"], default="normal", type=str) parser.add_argument("--savepath", type=str) return parser.parse_args() def main(): args = parse_args() # load annotations from json with open(os.path.join(args.path, f"scene_{args.scene:05d}", "annotation_3d.json")) as file: annos = json.load(file) if args.type == "wireframe": visualize_wireframe(annos) elif args.type == "plane": visualize_plane(annos, args) elif args.type == "floorplan": visualize_floorplan(annos) if __name__ == "__main__": main()