-
Notifications
You must be signed in to change notification settings - Fork 1
/
dataset.py
145 lines (115 loc) · 5.17 KB
/
dataset.py
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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import os.path as osp
import random
import cv2
import numpy as np
import paddle
from paddle.io import Dataset
from utils.misc import scandir
from utils.file_client import FileClient
from utils.img_utils import padding, img2tensor, imfrombytes
from data.transforms import paired_random_crop, random_augmentation
class Dataset_GaussianDenoising(Dataset):
"""Paired image dataset for image restoration.
Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and
GT image pairs.
There are three modes:
1. 'lmdb': Use lmdb files.
If opt['io_backend'] == lmdb.
2. 'meta_info_file': Use meta information file to generate paths.
If opt['io_backend'] != lmdb and opt['meta_info_file'] is not None.
3. 'folder': Scan folders to generate paths.
The rest.
Args:
opt (dict): Config for train datasets. It contains the following keys:
dataroot_gt (str): Data root path for gt.
meta_info_file (str): Path for meta information file.
io_backend (dict): IO backend type and other kwarg.
gt_size (int): Cropped patched size for gt patches.
use_flip (bool): Use horizontal flips.
use_rot (bool): Use rotation (use vertical flip and transposing h
and w for implementation).
scale (bool): Scale, which will be added automatically.
phase (str): 'train' or 'val'.
"""
def __init__(self, opt):
super(Dataset_GaussianDenoising, self).__init__()
self.opt = opt
if self.opt['phase'] == 'train':
self.sigma_type = opt['sigma_type']
self.sigma_range = opt['sigma_range']
assert self.sigma_type in ['constant', 'random', 'choice']
else:
self.sigma_test = opt['sigma_test']
self.in_ch = opt['in_ch']
# file client (io backend)
self.file_client = None
self.io_backend_opt = opt['io_backend']
self.mean = opt['mean'] if 'mean' in opt else None
self.std = opt['std'] if 'std' in opt else None
self.gt_folder = opt['dataroot_gt']
self.paths = sorted(list(scandir(self.gt_folder, full_path=True)))
if self.opt['phase'] == 'train':
self.geometric_augs = self.opt['geometric_augs']
def __getitem__(self, index):
if self.file_client is None:
self.file_client = FileClient(
self.io_backend_opt.pop('type'), **self.io_backend_opt)
scale = self.opt['scale']
index = index % len(self.paths)
# Load gt and lq images. Dimension order: HWC; channel order: BGR;
# image range: [0, 1], float32.
gt_path = self.paths[index]
img_bytes = self.file_client.get(gt_path, 'gt')
if self.in_ch == 3:
try:
img_gt = imfrombytes(img_bytes, float32=True)
except:
raise Exception("gt path {} not working".format(gt_path))
img_gt = cv2.cvtColor(img_gt, cv2.COLOR_BGR2RGB)
else:
try:
img_gt = imfrombytes(img_bytes, flag='grayscale', float32=True)
except:
raise Exception("gt path {} not working".format(gt_path))
img_gt = np.expand_dims(img_gt, axis=2)
img_lq = img_gt.copy()
# augmentation for training
if self.opt['phase'] == 'train':
gt_size = self.opt['gt_size']
# padding
img_gt, img_lq = padding(img_gt, img_lq, gt_size)
# random crop
img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale,
gt_path)
# flip, rotation
if self.geometric_augs:
img_gt, img_lq = random_augmentation(img_gt, img_lq)
img_gt, img_lq = img2tensor([img_gt, img_lq],
bgr2rgb=False,
float32=True)
if self.sigma_type == 'constant':
sigma_value = self.sigma_range
elif self.sigma_type == 'random':
sigma_value = random.uniform(self.sigma_range[0], self.sigma_range[1])
elif self.sigma_type == 'choice':
sigma_value = random.choice(self.sigma_range)
noise_level = sigma_value / 255.0
# noise_level_map = torch.ones((1, img_lq.size(1), img_lq.size(2))).mul_(noise_level).float()
# noise = np.random.randn(*img_lq.shape) * noise_level
noise = paddle.randn(img_lq.shape,dtype='float32').numpy() * noise_level
img_lq = img_lq + noise.astype('float32')
else:
np.random.seed(seed=0)
img_lq += np.random.normal(0, self.sigma_test/255.0, img_lq.shape)
# noise_level_map = torch.ones((1, img_lq.shape[0], img_lq.shape[1])).mul_(self.sigma_test/255.0).float()
img_gt, img_lq = img2tensor([img_gt, img_lq],
bgr2rgb=False,
float32=True)
return {
'lq': img_lq,
'gt': img_gt,
'lq_path': gt_path,
'gt_path': gt_path
}
def __len__(self):
return len(self.paths)