Spaces:
Configuration error
Configuration error
File size: 5,057 Bytes
89c5d90 |
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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import tensorflow as tf
from PIL import Image
from models.dexined import dexined
# from models.dexinedBs import dexined
from utls.utls import *
from utls.dataset_manager import data_parser,get_single_image,\
get_testing_batch
class m_tester():
def __init__(self, args):
self.args = args
def setup(self, session):
try:
if self.args.model_name=='DXN':
self.model = dexined(self.args)
else:
print_error("Error setting model, {}".format(self.args.model_name))
if self.args.trained_model_dir is None:
meta_model_file = os.path.join(
self.args.checkpoint_dir, os.path.join(
self.args.model_name + '_' + self.args.train_dataset,
os.path.join('train',
'{}-{}'.format(self.args.model_name, self.args.test_snapshot))))
else:
meta_model_file = os.path.join(
self.args.checkpoint_dir, os.path.join(
self.args.model_name + '_' + self.args.train_dataset,
os.path.join(self.args.trained_model_dir,
'{}-{}'.format(self.args.model_name, self.args.test_snapshot))))
saver = tf.train.Saver()
saver.restore(session, meta_model_file)
print_info('Done restoring DexiNed model from {}'.format(meta_model_file))
except Exception as err:
print_error('Error setting up DexiNed traied model, {}'.format(err))
def run(self, session):
self.model.setup_testing(session)
if self.args.use_dataset:
test_data= data_parser(self.args)
n_data = len(test_data[1])
else:
test_data=get_single_image(self.args)
n_data = len(test_data)
print_info('Writing PNGs at {}'.format(self.args.base_dir_results))
if self.args.batch_size_test==1 and self.args.use_dataset:
for i in range(n_data):
im, em, file_name = get_testing_batch(self.args,
[test_data[0][test_data[1][i]], test_data[1][i]], use_batch=False)
self.img_info = file_name
edgemap = session.run(self.model.predictions, feed_dict={self.model.images: [im]})
self.save_egdemaps(edgemap, single_image=True)
print_info('Done testing {}, {}'.format(self.img_info[0], self.img_info[1]))
# for individual images
elif self.args.batch_size_test==1 and not self.args.use_dataset:
for i in range(n_data):
im, file_name = get_single_image(self.args,file_path=test_data[i])
self.img_info = file_name
edgemap = session.run(self.model.predictions, feed_dict={self.model.images: [im]})
self.save_egdemaps(edgemap, single_image=True)
print_info('Done testing {}, {}'.format(self.img_info[0], self.img_info[1]))
def save_egdemaps(self, em_maps, single_image=False):
""" save_edgemaps descriptios
:param em_maps:
:param single_image:
save predicted edge maps
"""
result_dir = 'DexiNed_'+self.args.train_dataset+'2'+self.args.test_dataset
if self.args.base_dir_results is None:
res_dir = os.path.join('../results', result_dir)
else:
res_dir = os.path.join(self.args.base_dir_results,result_dir)
gt_dir = os.path.join(res_dir,'gt')
all_dir = os.path.join(res_dir,'pred-h5')
resf_dir = os.path.join(res_dir,'pred-f')
resa_dir = os.path.join(res_dir,'pred-a')
os.makedirs(resf_dir, exist_ok=True)
os.makedirs(resa_dir, exist_ok=True)
os.makedirs(gt_dir, exist_ok=True)
os.makedirs(all_dir, exist_ok=True)
if single_image:
em_maps = [e[0] for e in em_maps]
em_a = np.mean(np.array(em_maps), axis=0)
em_maps = em_maps + [em_a ]
em = em_maps[len(em_maps)-2]
em[em < self.args.testing_threshold] = 0.0
em_a[em_a < self.args.testing_threshold] = 0.0
em = 255.0 * (1.0 - em)
em_a = 255.0 * (1.0 - em_a)
em = np.tile(em, [1, 1, 3])
em_a = np.tile(em_a, [1, 1, 3])
em = Image.fromarray(np.uint8(em))
em_a = Image.fromarray(np.uint8(em_a))
tmp_name = os.path.basename(self.img_info[0])
tmp_name = tmp_name[:-4]
tmp_size = self.img_info[-1][:2]
tmp_size = (tmp_size[1],tmp_size[0])
em_f = em.resize(tmp_size)
em_a = em_a.resize(tmp_size)
em_f.save(os.path.join(resf_dir, tmp_name + '.png'))
em_a.save(os.path.join(resa_dir, tmp_name + '.png'))
em_maps =tensor_norm_01(em_maps)
save_variable_h5(os.path.join(all_dir, tmp_name + '.h5'), np.float16(em_maps))
else:
pass |