BasicSR

1
import math
2
import os
3
import torch
4
import torchvision.utils
5

6
from basicsr.data import build_dataloader, build_dataset
7

8

9
def main():
10
    """Test FFHQ dataset."""
11
    opt = {}
12
    opt['dist'] = False
13
    opt['gpu_ids'] = [0]
14
    opt['phase'] = 'train'
15

16
    opt['name'] = 'FFHQ'
17
    opt['type'] = 'FFHQDataset'
18

19
    opt['dataroot_gt'] = 'datasets/ffhq/ffhq_256.lmdb'
20
    opt['io_backend'] = dict(type='lmdb')
21

22
    opt['use_hflip'] = True
23
    opt['mean'] = [0.5, 0.5, 0.5]
24
    opt['std'] = [0.5, 0.5, 0.5]
25

26
    opt['num_worker_per_gpu'] = 1
27
    opt['batch_size_per_gpu'] = 4
28

29
    opt['dataset_enlarge_ratio'] = 1
30

31
    os.makedirs('tmp', exist_ok=True)
32

33
    dataset = build_dataset(opt)
34
    data_loader = build_dataloader(dataset, opt, num_gpu=0, dist=opt['dist'], sampler=None)
35

36
    nrow = int(math.sqrt(opt['batch_size_per_gpu']))
37
    padding = 2 if opt['phase'] == 'train' else 0
38

39
    print('start...')
40
    for i, data in enumerate(data_loader):
41
        if i > 5:
42
            break
43
        print(i)
44

45
        gt = data['gt']
46
        print(torch.min(gt), torch.max(gt))
47
        gt_path = data['gt_path']
48
        print(gt_path)
49
        torchvision.utils.save_image(
50
            gt, f'tmp/gt_{i:03d}.png', nrow=nrow, padding=padding, normalize=True, range=(-1, 1))
51

52

53
if __name__ == '__main__':
54
    main()
55