-
Notifications
You must be signed in to change notification settings - Fork 110
/
utils.py
222 lines (176 loc) · 8.29 KB
/
utils.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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
from PIL import Image
import os
import json
import random
import torchvision.transforms.functional as FT
import torch
import math
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Some constants
rgb_weights = torch.FloatTensor([65.481, 128.553, 24.966]).to(device)
imagenet_mean = torch.FloatTensor([0.485, 0.456, 0.406]).unsqueeze(1).unsqueeze(2)
imagenet_std = torch.FloatTensor([0.229, 0.224, 0.225]).unsqueeze(1).unsqueeze(2)
imagenet_mean_cuda = torch.FloatTensor([0.485, 0.456, 0.406]).to(device).unsqueeze(0).unsqueeze(2).unsqueeze(3)
imagenet_std_cuda = torch.FloatTensor([0.229, 0.224, 0.225]).to(device).unsqueeze(0).unsqueeze(2).unsqueeze(3)
def create_data_lists(train_folders, test_folders, min_size, output_folder):
"""
Create lists for images in the training set and each of the test sets.
:param train_folders: folders containing the training images; these will be merged
:param test_folders: folders containing the test images; each test folder will form its own test set
:param min_size: minimum width and height of images to be considered
:param output_folder: save data lists here
"""
print("\nCreating data lists... this may take some time.\n")
train_images = list()
for d in train_folders:
for i in os.listdir(d):
img_path = os.path.join(d, i)
img = Image.open(img_path, mode='r')
if img.width >= min_size and img.height >= min_size:
train_images.append(img_path)
print("There are %d images in the training data.\n" % len(train_images))
with open(os.path.join(output_folder, 'train_images.json'), 'w') as j:
json.dump(train_images, j)
for d in test_folders:
test_images = list()
test_name = d.split("/")[-1]
for i in os.listdir(d):
img_path = os.path.join(d, i)
img = Image.open(img_path, mode='r')
if img.width >= min_size and img.height >= min_size:
test_images.append(img_path)
print("There are %d images in the %s test data.\n" % (len(test_images), test_name))
with open(os.path.join(output_folder, test_name + '_test_images.json'), 'w') as j:
json.dump(test_images, j)
print("JSONS containing lists of Train and Test images have been saved to %s\n" % output_folder)
def convert_image(img, source, target):
"""
Convert an image from a source format to a target format.
:param img: image
:param source: source format, one of 'pil' (PIL image), '[0, 1]' or '[-1, 1]' (pixel value ranges)
:param target: target format, one of 'pil' (PIL image), '[0, 255]', '[0, 1]', '[-1, 1]' (pixel value ranges),
'imagenet-norm' (pixel values standardized by imagenet mean and std.),
'y-channel' (luminance channel Y in the YCbCr color format, used to calculate PSNR and SSIM)
:return: converted image
"""
assert source in {'pil', '[0, 1]', '[-1, 1]'}, "Cannot convert from source format %s!" % source
assert target in {'pil', '[0, 255]', '[0, 1]', '[-1, 1]', 'imagenet-norm',
'y-channel'}, "Cannot convert to target format %s!" % target
# Convert from source to [0, 1]
if source == 'pil':
img = FT.to_tensor(img)
elif source == '[0, 1]':
pass # already in [0, 1]
elif source == '[-1, 1]':
img = (img + 1.) / 2.
# Convert from [0, 1] to target
if target == 'pil':
img = FT.to_pil_image(img)
elif target == '[0, 255]':
img = 255. * img
elif target == '[0, 1]':
pass # already in [0, 1]
elif target == '[-1, 1]':
img = 2. * img - 1.
elif target == 'imagenet-norm':
if img.ndimension() == 3:
img = (img - imagenet_mean) / imagenet_std
elif img.ndimension() == 4:
img = (img - imagenet_mean_cuda) / imagenet_std_cuda
elif target == 'y-channel':
# Based on definitions at https://github.com/xinntao/BasicSR/wiki/Color-conversion-in-SR
# torch.dot() does not work the same way as numpy.dot()
# So, use torch.matmul() to find the dot product between the last dimension of an 4-D tensor and a 1-D tensor
img = torch.matmul(255. * img.permute(0, 2, 3, 1)[:, 4:-4, 4:-4, :], rgb_weights) / 255. + 16.
return img
class ImageTransforms(object):
"""
Image transformation pipeline.
"""
def __init__(self, split, crop_size, scaling_factor, lr_img_type, hr_img_type):
"""
:param split: one of 'train' or 'test'
:param crop_size: crop size of HR images
:param scaling_factor: LR images will be downsampled from the HR images by this factor
:param lr_img_type: the target format for the LR image; see convert_image() above for available formats
:param hr_img_type: the target format for the HR image; see convert_image() above for available formats
"""
self.split = split.lower()
self.crop_size = crop_size
self.scaling_factor = scaling_factor
self.lr_img_type = lr_img_type
self.hr_img_type = hr_img_type
assert self.split in {'train', 'test'}
def __call__(self, img):
"""
:param img: a PIL source image from which the HR image will be cropped, and then downsampled to create the LR image
:return: LR and HR images in the specified format
"""
# Crop
if self.split == 'train':
# Take a random fixed-size crop of the image, which will serve as the high-resolution (HR) image
left = random.randint(1, img.width - self.crop_size)
top = random.randint(1, img.height - self.crop_size)
right = left + self.crop_size
bottom = top + self.crop_size
hr_img = img.crop((left, top, right, bottom))
else:
# Take the largest possible center-crop of it such that its dimensions are perfectly divisible by the scaling factor
x_remainder = img.width % self.scaling_factor
y_remainder = img.height % self.scaling_factor
left = x_remainder // 2
top = y_remainder // 2
right = left + (img.width - x_remainder)
bottom = top + (img.height - y_remainder)
hr_img = img.crop((left, top, right, bottom))
# Downsize this crop to obtain a low-resolution version of it
lr_img = hr_img.resize((int(hr_img.width / self.scaling_factor), int(hr_img.height / self.scaling_factor)),
Image.BICUBIC)
# Sanity check
assert hr_img.width == lr_img.width * self.scaling_factor and hr_img.height == lr_img.height * self.scaling_factor
# Convert the LR and HR image to the required type
lr_img = convert_image(lr_img, source='pil', target=self.lr_img_type)
hr_img = convert_image(hr_img, source='pil', target=self.hr_img_type)
return lr_img, hr_img
class AverageMeter(object):
"""
Keeps track of most recent, average, sum, and count of a metric.
"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def clip_gradient(optimizer, grad_clip):
"""
Clips gradients computed during backpropagation to avoid explosion of gradients.
:param optimizer: optimizer with the gradients to be clipped
:param grad_clip: clip value
"""
for group in optimizer.param_groups:
for param in group['params']:
if param.grad is not None:
param.grad.data.clamp_(-grad_clip, grad_clip)
def save_checkpoint(state, filename):
"""
Save model checkpoint.
:param state: checkpoint contents
"""
torch.save(state, filename)
def adjust_learning_rate(optimizer, shrink_factor):
"""
Shrinks learning rate by a specified factor.
:param optimizer: optimizer whose learning rate must be shrunk.
:param shrink_factor: factor in interval (0, 1) to multiply learning rate with.
"""
print("\nDECAYING learning rate.")
for param_group in optimizer.param_groups:
param_group['lr'] = param_group['lr'] * shrink_factor
print("The new learning rate is %f\n" % (optimizer.param_groups[0]['lr'],))