File size: 3,038 Bytes
424188c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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()
|