| 23 | |
| 24 | |
| 25 | def create_dataset(opt,dataset_opt): |
| 26 | mode = dataset_opt['mode'] |
| 27 | # datasets for image restoration |
| 28 | if mode == 'UEN_train': |
| 29 | from data.SIEN_dataset import DatasetFromFolder as D |
| 30 | |
| 31 | dataset = D(upscale_factor=opt['scale'], data_augmentation=dataset_opt['augment'], |
| 32 | group_file=dataset_opt['filelist'], |
| 33 | patch_size=dataset_opt['IN_size'], black_edges_crop=False, hflip=True, rot=True) |
| 34 | |
| 35 | elif mode == 'UEN_val': |
| 36 | from data.SIEN_dataset import DatasetFromFolder as D |
| 37 | print('########### val loader with pacth:', dataset_opt['IN_size']) |
| 38 | dataset = D(upscale_factor=opt['scale'], data_augmentation=False, |
| 39 | group_file=dataset_opt['filelist'], |
| 40 | ################## |
| 41 | # 2021124 |
| 42 | # validation 直接用原图太大会爆,改用patch |
| 43 | patch_size=256, black_edges_crop=False, hflip=False, rot=False) |
| 44 | # patch_size=None, black_edges_crop=False, hflip=False, rot=False) |
| 45 | ############### |
| 46 | # 20221205 |
| 47 | elif mode == 'UEN_test': |
| 48 | from data.SIEN_dataset import DatasetFromFolder as D |
| 49 | print('########### val loader with pacth:', dataset_opt['IN_size']) |
| 50 | dataset = D(upscale_factor=opt['scale'], data_augmentation=False, |
| 51 | group_file=dataset_opt['filelist'], |
| 52 | ################## |
| 53 | # 20221124 |
| 54 | # validation 直接用原图太大会爆,改用patch |
| 55 | patch_size=None |
| 56 | , black_edges_crop=False, hflip=False, rot=False) |
| 57 | ############### |
| 58 | ############### |
| 59 | # 20230328 |
| 60 | elif mode == 'multi_task': |
| 61 | from data.multitask_dataset import DatasetFromFolder as D |
| 62 | |
| 63 | dataset = D(upscale_factor=opt['scale'], data_augmentation=dataset_opt['augment'], |
| 64 | group_file=dataset_opt['filelist'], |
| 65 | patch_size=dataset_opt['IN_size'], black_edges_crop=False, hflip=True, rot=True) |
| 66 | ############## |
| 67 | |
| 68 | else: |
| 69 | raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode)) |
| 70 | |
| 71 | |
| 72 | logger = logging.getLogger('base') |
| 73 | logger.info('Dataset [{:s} - {:s}] is created.'.format(dataset.__class__.__name__, |
| 74 | dataset_opt['name'])) |
| 75 | return dataset |