mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow test to operate on batches

Browse Source
This commit is contained in:
James Betker 2020-04-23 23:59:09 -06:00
parent 8ead9ae183
commit e98d92fc77
2 changed files with 92 additions and 83 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
codes/data/__init__.py
View File

@ -21,7 +21,7 @@ def create_dataloader(dataset, dataset_opt, opt=None, sampler=None):
num_workers=num_workers, sampler=sampler, drop_last=True, num_workers=num_workers, sampler=sampler, drop_last=True,
pin_memory=False) pin_memory=False)
else: else:
return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0, return torch.utils.data.DataLoader(dataset, batch_size=12, shuffle=False, num_workers=3,
pin_memory=False) pin_memory=False)
@ -32,8 +32,9 @@ def create_dataset(dataset_opt):
from data.LQ_dataset import LQDataset as D from data.LQ_dataset import LQDataset as D
elif mode == 'LQGT': elif mode == 'LQGT':
from data.LQGT_dataset import LQGTDataset as D from data.LQGT_dataset import LQGTDataset as D
elif mode == 'GTLQ': # datasets for image corruption
from data.GTLQ_dataset import GTLQDataset as D elif mode == 'downsample':
from data.Downsample_dataset import DownsampleDataset as D
# datasets for video restoration # datasets for video restoration
elif mode == 'REDS': elif mode == 'REDS':
from data.REDS_dataset import REDSDataset as D from data.REDS_dataset import REDSDataset as D

24
codes/test.py
View File

@ -9,10 +9,13 @@ import utils.util as util
from data.util import bgr2ycbcr from data.util import bgr2ycbcr
from data import create_dataset, create_dataloader from data import create_dataset, create_dataloader
from models import create_model from models import create_model
from tqdm import tqdm
if __name__ == "__main__":
#### options #### options
want_just_images = True
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, help='Path to options YMAL file.', default='options/test/test_ESRGAN_vrp.yml') parser.add_argument('-opt', type=str, help='Path to options YMAL file.', default='options/test/test_corrupt_vixen_adrianna.yml')
opt = option.parse(parser.parse_args().opt, is_train=False) opt = option.parse(parser.parse_args().opt, is_train=False)
opt = option.dict_to_nonedict(opt) opt = option.dict_to_nonedict(opt)
@ -46,16 +49,18 @@ for test_loader in test_loaders:
test_results['psnr_y'] = [] test_results['psnr_y'] = []
test_results['ssim_y'] = [] test_results['ssim_y'] = []
for data in test_loader: tq = tqdm(test_loader)
for data in tq:
need_GT = False if test_loader.dataset.opt['dataroot_GT'] is None else True need_GT = False if test_loader.dataset.opt['dataroot_GT'] is None else True
model.feed_data(data, need_GT=need_GT) model.feed_data(data, need_GT=need_GT)
img_path = data['GT_path'][0] if need_GT else data['LQ_path'][0] model.test()
visuals = model.fake_H.detach().float().cpu()
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] img_name = osp.splitext(osp.basename(img_path))[0]
model.test() sr_img = util.tensor2img(visuals[i]) # uint8
visuals = model.get_current_visuals(need_GT=need_GT)
sr_img = util.tensor2img(visuals['rlt']) # uint8
# save images # save images
suffix = opt['suffix'] suffix = opt['suffix']
@ -65,6 +70,9 @@ for test_loader in test_loaders:
save_img_path = osp.join(dataset_dir, img_name + '.png') save_img_path = osp.join(dataset_dir, img_name + '.png')
util.save_img(sr_img, save_img_path) util.save_img(sr_img, save_img_path)
if want_just_images:
continue
# calculate PSNR and SSIM # calculate PSNR and SSIM
if need_GT: if need_GT:
gt_img = util.tensor2img(visuals['GT']) gt_img = util.tensor2img(visuals['GT'])
@ -90,7 +98,7 @@ for test_loader in test_loaders:
else: else:
logger.info(img_name) logger.info(img_name)
if need_GT: # metrics if not want_just_images and need_GT: # metrics
# Average PSNR/SSIM results # Average PSNR/SSIM results
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])