MAG-ViT-Super-Resolution/codes/data/__init__.py at main · riotu-lab/MAG-ViT-Super-Resolution · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

from torch.utils.data import DataLoader

def create_dataloaders(args):
    """create dataloader"""
    if args.dataset == 'AID':
        from data.aid import AIDataset
        training_set = AIDataset(args, root_dir='/data/Image_restoration/Datasets/AID-dataset/train',  # AID_dataset
                                 train=True)
        val_set = AIDataset(args, root_dir='/data/Image_restoration/Datasets/AID-dataset/val',
                            train=False)
    elif args.dataset == 'UCMerced':
        from data.ucmerced import UCMercedDataset
        training_set = UCMercedDataset(args, root_dir='/data/Image_restoration/Datasets/UCMerced-dataset/train',
                                 train=True)
        val_set = UCMercedDataset(args, root_dir='/data/Image_restoration/Datasets/UCMerced-dataset/val',
                            train=False)
    elif args.dataset == 'DIV2K':
        from data.div2k import DIV2KDataset
        training_set = DIV2KDataset(args, root_dir='F:/research/dataset/SR for natural image/DIV2K_train',
                                    train=True)
        val_set = DIV2KDataset(args, root_dir='F:/research/dataset/SR for natural image/valid',
                               train=False)
    else:
        raise NotImplementedError(
            'Wrong dataset name %s ' % args.dataset)

    dataloaders = {'train': DataLoader(training_set, batch_size=args.batch_size,
                                 shuffle=True, num_workers=0),  # args.n_threads
                   'val': DataLoader(val_set, batch_size=args.batch_size,
                                 shuffle=True, num_workers=0)}  # args.n_threads

    return dataloaders