def visdom_image(img_dict, window):
for idx, key in enumerate(img_dict):
win = window + idx
tensor_img = train_utils.tensor2im(img_dict[key].data)
visdom.image(tensor_img.transpose([2, 0, 1]),
opts=dict(title=key),
win=win)
def test(loader_test, visualAttentionNet, root_dir):
visualAttentionNet.eval()
for itr, data in enumerate(loader_test):
testImg, fileName = data[0], data[1]
testImg = testImg.cuda()
with torch.no_grad():
test_attention_result = visualAttentionNet(testImg)
test_recon_result_img = train_utils.tensor2im(
test_attention_result)
norm_input_img = train_utils.tensor2im(testImg +
test_attention_result)
recon_save_dir = root_dir + 'visual_attention_map_' + fileName[
0].split('.')[0] + ('.png')
recon_save_dir2 = root_dir + 'sum_' + fileName[0].split('.')[0] + (
'.png')
train_utils.save_images(test_recon_result_img, recon_save_dir)
train_utils.save_images(norm_input_img, recon_save_dir2)
def test(args, loader_test, model_AttentionNet, epoch, root_dir) :
model_AttentionNet.eval()
for itr, data in enumerate(loader_test):
testImg, fileName = data[0], data[1]
if args.cuda:
testImg = testImg.cuda()
with torch.no_grad():
test_result = model_AttentionNet(testImg)
test_result_img = train_utils.tensor2im(test_result)
result_save_dir = root_dir + fileName[0].split('.')[0]+('_epoch_{}_itr_{}.png'.format(epoch, itr))
train_utils.save_images(test_result_img, result_save_dir)
def test(loader_test, VAN, EN, root_dir):
VAN.eval()
EN.eval()
for itr, data in enumerate(loader_test):
testImg, img_name = data[0], data[1]
testImg = testImg.cuda()
with torch.no_grad():
visual_attention_map = VAN(testImg)
enhance_result = EN(testImg, visual_attention_map)
enhance_result_img = train_utils.tensor2im(enhance_result)
result_save_dir = root_dir + 'enhance'+ img_name[0].split('.')[0]+('.png')
train_utils.save_images(enhance_result_img, result_save_dir)