File size: 1,116 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 |
import os
import numpy as np
import cv2
from planar_graph_utils import get_regions_from_pg, plot_floorplan_with_regions
# example_pg = {
# (127, 20): [(20, 120), (234, 120)],
# (20, 120): [(127, 20), (234, 120), (20, 240)],
# (234, 120): [(127, 20), (20, 120), (234, 240)],
# (20, 240): [(20, 120), (234, 240)],
# (234, 240): [(234, 120), (20, 240)],
# }
pg_base = '../results/npy_heat_s3d_256/'
viz_base = './viz_heat'
if not os.path.exists(viz_base):
os.makedirs(viz_base)
for filename in sorted(os.listdir(pg_base)):
pg_path = os.path.join(pg_base, filename)
example_pg = np.load(pg_path, allow_pickle=True).tolist()
corners = example_pg['corners']
corners = corners.astype(np.int)
edges = example_pg['edges']
print('Processing file: {}'.format(filename))
regions = get_regions_from_pg(example_pg, corner_sorted=True)
print('num of extracted regions {}'.format(len(regions)))
floorplan_image = plot_floorplan_with_regions(regions, corners, edges, scale=1000)
cv2.imwrite(os.path.join(viz_base, '{}.png'.format(filename[:-4])), floorplan_image)
|