Spaces:
Runtime error
Runtime error
File size: 2,159 Bytes
419718d |
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 |
import argparse
import os
import sys
import glob
import time
import random
import numpy as np
import imageio
from PIL import Image
from matplotlib import pyplot as plt
sys.path.append('./utils/')
from rgb_ind_convertor import *
from util import *
parser = argparse.ArgumentParser()
parser.add_argument('--result_dir', type=str, default='./out',
help='The folder that save network predictions.')
def post_process(input_dir, save_dir, merge=True):
if not os.path.exists(save_dir):
os.mkdir(save_dir)
input_paths = sorted(glob.glob(os.path.join(input_dir, '*.png')))
names = [i.split('/')[-1] for i in input_paths]
out_paths = [os.path.join(save_dir, i) for i in names]
n = len(input_paths)
# n = 1
for i in range(n):
im = imageio.imread(input_paths[i], mode='RGB')
im_ind = rgb2ind(im, color_map=floorplan_fuse_map)
# seperate image into room-seg & boundary-seg
rm_ind = im_ind.copy()
rm_ind[im_ind==9] = 0
rm_ind[im_ind==10] = 0
bd_ind = np.zeros(im_ind.shape, dtype=np.uint8)
bd_ind[im_ind==9] = 9
bd_ind[im_ind==10] = 10
hard_c = (bd_ind>0).astype(np.uint8)
# region from room prediction it self
rm_mask = np.zeros(rm_ind.shape)
rm_mask[rm_ind>0] = 1
# region from close_wall line
cw_mask = hard_c
# refine close wall mask by filling the grap between bright line
cw_mask = fill_break_line(cw_mask)
fuse_mask = cw_mask + rm_mask
fuse_mask[fuse_mask>=1] = 255
# refine fuse mask by filling the hole
fuse_mask = flood_fill(fuse_mask)
fuse_mask = fuse_mask // 255
# one room one label
new_rm_ind = refine_room_region(cw_mask, rm_ind)
# ignore the background mislabeling
new_rm_ind = fuse_mask*new_rm_ind
print('Saving {}th refined room prediciton to {}'.format(i, out_paths[i]))
if merge:
new_rm_ind[bd_ind==9] = 9
new_rm_ind[bd_ind==10] = 10
rgb = ind2rgb(new_rm_ind, color_map=floorplan_fuse_map)
else:
rgb = ind2rgb(new_rm_ind)
imageio.imwrite(out_paths[i], rgb)
if __name__ == '__main__':
FLAGS, unparsed = parser.parse_known_args()
input_dir = FLAGS.result_dir
save_dir = os.path.join(input_dir, 'post')
post_process(input_dir, save_dir)
|