| import os | |
| import json | |
| import argparse | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from shapely.geometry import Polygon | |
| from descartes.patch import PolygonPatch | |
| from misc.panorama import draw_boundary_from_cor_id | |
| from misc.colors import colormap_255 | |
| def visualize_panorama(args): | |
| """visualize panorama layout | |
| """ | |
| scene_path = os.path.join(args.path, f"scene_{args.scene:05d}", "2D_rendering") | |
| for room_id in np.sort(os.listdir(scene_path)): | |
| room_path = os.path.join(scene_path, room_id, "panorama") | |
| cor_id = np.loadtxt(os.path.join(room_path, "layout.txt")) | |
| img_src = cv2.imread(os.path.join(room_path, "full", "rgb_rawlight.png")) | |
| img_src = cv2.cvtColor(img_src, cv2.COLOR_BGR2RGB) | |
| img_viz = draw_boundary_from_cor_id(cor_id, img_src) | |
| plt.axis('off') | |
| plt.imshow(img_viz) | |
| plt.show() | |
| def visualize_perspective(args): | |
| """visualize perspective layout | |
| """ | |
| colors = np.array(colormap_255) / 255 | |
| scene_path = os.path.join(args.path, f"scene_{args.scene:05d}", "2D_rendering") | |
| for room_id in np.sort(os.listdir(scene_path)): | |
| room_path = os.path.join(scene_path, room_id, "perspective", "full") | |
| if not os.path.exists(room_path): | |
| continue | |
| for position_id in np.sort(os.listdir(room_path)): | |
| position_path = os.path.join(room_path, position_id) | |
| image = cv2.imread(os.path.join(position_path, "rgb_rawlight.png")) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| with open(os.path.join(position_path, "layout.json")) as f: | |
| annos = json.load(f) | |
| fig = plt.figure() | |
| for i, key in enumerate(['amodal_mask', 'visible_mask']): | |
| ax = fig.add_subplot(2, 1, i + 1) | |
| plt.axis('off') | |
| plt.imshow(image) | |
| for i, planes in enumerate(annos['planes']): | |
| if len(planes[key]): | |
| for plane in planes[key]: | |
| polygon = Polygon([annos['junctions'][id]['coordinate'] for id in plane]) | |
| patch = PolygonPatch(polygon, facecolor=colors[i], alpha=0.5) | |
| ax.add_patch(patch) | |
| plt.title(key) | |
| plt.show() | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Structured3D 2D Layout 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=["perspective", "panorama"], required=True, | |
| help="type of camera", type=str) | |
| return parser.parse_args() | |
| def main(): | |
| args = parse_args() | |
| if args.type == 'panorama': | |
| visualize_panorama(args) | |
| elif args.type == 'perspective': | |
| visualize_perspective(args) | |
| if __name__ == "__main__": | |
| main() | |