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