|
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) |
|
|
|
|
|
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) |
|
|
|
|
|
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) |
|
|
|
|
|
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) |
|
|
|
|
|
line_colors = np.repeat(colormap[5].reshape(1, 3), len(junction_pairs), axis=0) |
|
|
|
|
|
if len(lines_holes) != 0: |
|
line_colors[lines_holes] = colormap[6] |
|
|
|
|
|
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 |
|
""" |
|
|
|
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) |
|
|
|
|
|
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'] |
|
|
|
|
|
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) |
|
|
|
|
|
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): |
|
|
|
|
|
|
|
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': |
|
|
|
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 |
|
""" |
|
|
|
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']) |
|
|
|
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() |
|
|
|
|
|
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() |
|
|
