114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
import argparse
|
|
import logging
|
|
import os.path as osp
|
|
import random
|
|
import time
|
|
from collections import OrderedDict
|
|
|
|
import numpy as np
|
|
import torch
|
|
from tqdm import tqdm
|
|
|
|
import dlas.utils
|
|
import dlas.utils.options as option
|
|
import dlas.utils.util as util
|
|
from dlas.data import create_dataloader, create_dataset
|
|
from dlas.trainer.ExtensibleTrainer import ExtensibleTrainer
|
|
|
|
|
|
def forward_pass(model, data, output_dir, opt):
|
|
alteration_suffix = util.opt_get(opt, ['name'], '')
|
|
denorm_range = tuple(util.opt_get(
|
|
opt, ['image_normalization_range'], [0, 1]))
|
|
with torch.no_grad():
|
|
model.feed_data(data, 0, need_GT=need_GT)
|
|
model.test()
|
|
|
|
visuals = model.get_current_visuals(need_GT)['rlt'].cpu()
|
|
visuals = (visuals - denorm_range[0]) / (denorm_range[1]-denorm_range[0])
|
|
fea_loss = 0
|
|
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:
|
|
psnr_sr = util.tensor2img(visuals[i])
|
|
psnr_gt = util.tensor2img(data['hq'][i])
|
|
psnr_loss += util.calculate_psnr(psnr_sr, psnr_gt)
|
|
|
|
util.save_img(sr_img, save_img_path)
|
|
return fea_loss, psnr_loss
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Set seeds
|
|
torch.manual_seed(5555)
|
|
random.seed(5555)
|
|
np.random.seed(5555)
|
|
|
|
# options
|
|
torch.backends.cudnn.benchmark = True
|
|
want_metrics = False
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('-opt', type=str, help='Path to options YAML file.',
|
|
default='../options/test_diffusion_unet.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
|
|
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
|
|
need_GT = need_GT and want_metrics
|
|
|
|
fea_loss, psnr_loss = forward_pass(model, data, dataset_dir, opt)
|
|
fea_loss += fea_loss
|
|
psnr_loss += psnr_loss
|
|
|
|
# log
|
|
logger.info('# Validation # Fea: {:.4e}, PSNR: {:.4e}'.format(
|
|
fea_loss / len(test_loader), psnr_loss / len(test_loader)))
|