| 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() |
|
|
