-
Notifications
You must be signed in to change notification settings - Fork 37
/
dataset.py
151 lines (114 loc) · 5.03 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
146
147
148
149
150
151
import torch.utils.data as data
from torchvision.transforms import *
from os import listdir
from os.path import join
from PIL import Image
import random
def is_image_file(filename):
return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".bmp"])
def load_img(filepath):
img = Image.open(filepath).convert('RGB')
return img
def calculate_valid_crop_size(crop_size, scale_factor):
return crop_size - (crop_size % scale_factor)
class TrainDatasetFromFolder(data.Dataset):
def __init__(self, image_dirs, is_gray=False, random_scale=True, crop_size=128, rotate=True, fliplr=True,
fliptb=True, scale_factor=4):
super(TrainDatasetFromFolder, self).__init__()
self.image_filenames = []
for image_dir in image_dirs:
self.image_filenames.extend(join(image_dir, x) for x in sorted(listdir(image_dir)) if is_image_file(x))
self.is_gray = is_gray
self.random_scale = random_scale
self.crop_size = crop_size
self.rotate = rotate
self.fliplr = fliplr
self.fliptb = fliptb
self.scale_factor = scale_factor
def __getitem__(self, index):
# load image
img = load_img(self.image_filenames[index])
# determine valid HR image size with scale factor
self.crop_size = calculate_valid_crop_size(self.crop_size, self.scale_factor)
hr_img_w = self.crop_size
hr_img_h = self.crop_size
# determine LR image size
lr_img_w = hr_img_w // self.scale_factor
lr_img_h = hr_img_h // self.scale_factor
# random scaling between [0.5, 1.0]
if self.random_scale:
eps = 1e-3
ratio = random.randint(5, 10) * 0.1
if hr_img_w * ratio < self.crop_size:
ratio = self.crop_size / hr_img_w + eps
if hr_img_h * ratio < self.crop_size:
ratio = self.crop_size / hr_img_h + eps
scale_w = int(hr_img_w * ratio)
scale_h = int(hr_img_h * ratio)
transform = Scale((scale_w, scale_h), interpolation=Image.BICUBIC)
img = transform(img)
# random crop
transform = RandomCrop(self.crop_size)
img = transform(img)
# random rotation between [90, 180, 270] degrees
if self.rotate:
rv = random.randint(1, 3)
img = img.rotate(90 * rv, expand=True)
# random horizontal flip
if self.fliplr:
transform = RandomHorizontalFlip()
img = transform(img)
# random vertical flip
if self.fliptb:
if random.random() < 0.5:
img = img.transpose(Image.FLIP_TOP_BOTTOM)
# only Y-channel is super-resolved
if self.is_gray:
img = img.convert('YCbCr')
# img, _, _ = img.split()
# hr_img HR image
hr_transform = Compose([Scale((hr_img_w, hr_img_h), interpolation=Image.BICUBIC), ToTensor()])
hr_img = hr_transform(img)
# lr_img LR image
lr_transform = Compose([Scale((lr_img_w, lr_img_h), interpolation=Image.BICUBIC), ToTensor()])
lr_img = lr_transform(img)
# Bicubic interpolated image
bc_transform = Compose([ToPILImage(), Scale((hr_img_w, hr_img_h), interpolation=Image.BICUBIC), ToTensor()])
bc_img = bc_transform(lr_img)
return lr_img, hr_img, bc_img
def __len__(self):
return len(self.image_filenames)
class TestDatasetFromFolder(data.Dataset):
def __init__(self, image_dir, is_gray=False, scale_factor=4):
super(TestDatasetFromFolder, self).__init__()
self.image_filenames = [join(image_dir, x) for x in sorted(listdir(image_dir)) if is_image_file(x)]
self.is_gray = is_gray
self.scale_factor = scale_factor
def __getitem__(self, index):
# load image
img = load_img(self.image_filenames[index])
# original HR image size
w = img.size[0]
h = img.size[1]
# determine valid HR image size with scale factor
hr_img_w = calculate_valid_crop_size(w, self.scale_factor)
hr_img_h = calculate_valid_crop_size(h, self.scale_factor)
# determine lr_img LR image size
lr_img_w = hr_img_w // self.scale_factor
lr_img_h = hr_img_h // self.scale_factor
# only Y-channel is super-resolved
if self.is_gray:
img = img.convert('YCbCr')
# img, _, _ = lr_img.split()
# hr_img HR image
hr_transform = Compose([Scale((hr_img_w, hr_img_h), interpolation=Image.BICUBIC), ToTensor()])
hr_img = hr_transform(img)
# lr_img LR image
lr_transform = Compose([Scale((lr_img_w, lr_img_h), interpolation=Image.BICUBIC), ToTensor()])
lr_img = lr_transform(img)
# Bicubic interpolated image
bc_transform = Compose([ToPILImage(), Scale((hr_img_w, hr_img_h), interpolation=Image.BICUBIC), ToTensor()])
bc_img = bc_transform(lr_img)
return lr_img, hr_img, bc_img
def __len__(self):
return len(self.image_filenames)