DL-Art-School/codes/test.py

99 lines
3.5 KiB
Python
Raw Normal View History

import os.path as osp
import logging
import time
import argparse
from collections import OrderedDict
import utils
import utils.options as option
import utils.util as util
2020-12-18 16:18:34 +00:00
from trainer.ExtensibleTrainer import ExtensibleTrainer
from data import create_dataset, create_dataloader
from tqdm import tqdm
import torch
def forward_pass(model, output_dir, alteration_suffix=''):
2020-11-29 19:21:31 +00:00
model.feed_data(data, 0, need_GT=need_GT)
model.test()
visuals = model.get_current_visuals(need_GT)['rlt'].cpu()
fea_loss = 0
2020-10-27 16:25:42 +00:00
psnr_loss = 0
for i in range(visuals.shape[0]):
img_path = data['GT_path'][i] if need_GT else data['LQ_path'][i]
img_name = osp.splitext(osp.basename(img_path))[0]
sr_img = util.tensor2img(visuals[i]) # uint8
# save images
suffix = alteration_suffix
if suffix:
save_img_path = osp.join(output_dir, img_name + suffix + '.png')
else:
save_img_path = osp.join(output_dir, img_name + '.png')
if need_GT:
fea_loss += model.compute_fea_loss(visuals[i], data['hq'][i])
2020-10-27 16:25:42 +00:00
psnr_sr = util.tensor2img(visuals[i])
psnr_gt = util.tensor2img(data['hq'][i])
2020-10-27 16:25:42 +00:00
psnr_loss += util.calculate_psnr(psnr_sr, psnr_gt)
util.save_img(sr_img, save_img_path)
2020-10-27 16:25:42 +00:00
return fea_loss, psnr_loss
2020-07-24 18:26:44 +00:00
2020-04-24 05:59:09 +00:00
if __name__ == "__main__":
#### options
torch.backends.cudnn.benchmark = True
2020-11-29 19:21:31 +00:00
want_metrics = False
parser = argparse.ArgumentParser()
2020-12-22 22:42:38 +00:00
parser.add_argument('-opt', type=str, help='Path to options YAML file.', default='../options/test_faces_glean.yml')
opt = option.parse(parser.parse_args().opt, is_train=False)
opt = option.dict_to_nonedict(opt)
utils.util.loaded_options = opt
util.mkdirs(
(path for key, path in opt['path'].items()
if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key))
util.setup_logger('base', opt['path']['log'], 'test_' + opt['name'], level=logging.INFO,
screen=True, tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt))
#### Create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set, dataset_opt)
logger.info('Number of test images in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set)))
test_loaders.append(test_loader)
model = ExtensibleTrainer(opt)
fea_loss = 0
2020-10-27 16:25:42 +00:00
psnr_loss = 0
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
logger.info('\nTesting [{:s}]...'.format(test_set_name))
test_start_time = time.time()
dataset_dir = osp.join(opt['path']['results_root'], test_set_name)
util.mkdir(dataset_dir)
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
test_results['psnr_y'] = []
test_results['ssim_y'] = []
tq = tqdm(test_loader)
for data in tq:
need_GT = False if test_loader.dataset.opt['dataroot_GT'] is None else True
2020-11-29 19:21:31 +00:00
need_GT = need_GT and want_metrics
2020-11-27 03:30:55 +00:00
fea_loss, psnr_loss = forward_pass(model, dataset_dir, opt['name'])
fea_loss += fea_loss
psnr_loss += psnr_loss
# log
2020-10-27 16:25:42 +00:00
logger.info('# Validation # Fea: {:.4e}, PSNR: {:.4e}'.format(fea_loss / len(test_loader), psnr_loss / len(test_loader)))