AlexZou commited on
Commit
9eb95ee
1 Parent(s): 8b80ab2

Upload 7 files

Browse files
utils/__pycache__/utils_image.cpython-36.pyc ADDED
Binary file (18 kB). View file
 
utils/__pycache__/utils_image.cpython-37.pyc ADDED
Binary file (17.7 kB). View file
 
utils/__pycache__/utils_logger.cpython-36.pyc ADDED
Binary file (1.67 kB). View file
 
utils/__pycache__/utils_logger.cpython-37.pyc ADDED
Binary file (1.69 kB). View file
 
utils/test.bmp ADDED
utils/utils_image.py ADDED
@@ -0,0 +1,778 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import math
3
+ import random
4
+ import numpy as np
5
+ import torch
6
+ import cv2
7
+ from torchvision.utils import make_grid
8
+ from datetime import datetime
9
+ # import torchvision.transforms as transforms
10
+ import matplotlib.pyplot as plt
11
+
12
+ '''
13
+ modified by Kai Zhang (github: https://github.com/cszn)
14
+ 03/03/2019
15
+ https://github.com/twhui/SRGAN-pyTorch
16
+ https://github.com/xinntao/BasicSR
17
+ '''
18
+
19
+ IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP']
20
+
21
+
22
+ def is_image_file(filename):
23
+ return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
24
+
25
+
26
+ def get_timestamp():
27
+ return datetime.now().strftime('%y%m%d-%H%M%S')
28
+
29
+
30
+ def imshow(x, title=None, cbar=False, figsize=None):
31
+ plt.figure(figsize=figsize)
32
+ plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')
33
+ if title:
34
+ plt.title(title)
35
+ if cbar:
36
+ plt.colorbar()
37
+ plt.show()
38
+
39
+
40
+ '''
41
+ # =======================================
42
+ # get image pathes of files
43
+ # =======================================
44
+ '''
45
+
46
+
47
+ def get_image_paths(dataroot):
48
+ paths = None # return None if dataroot is None
49
+ if dataroot is not None:
50
+ paths = sorted(_get_paths_from_images(dataroot))
51
+ return paths
52
+
53
+
54
+ def _get_paths_from_images(path):
55
+ assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
56
+ images = []
57
+ for dirpath, _, fnames in sorted(os.walk(path)):
58
+ for fname in sorted(fnames):
59
+ if is_image_file(fname):
60
+ img_path = os.path.join(dirpath, fname)
61
+ images.append(img_path)
62
+ assert images, '{:s} has no valid image file'.format(path)
63
+ return images
64
+
65
+
66
+ '''
67
+ # =======================================
68
+ # makedir
69
+ # =======================================
70
+ '''
71
+
72
+
73
+ def mkdir(path):
74
+ if not os.path.exists(path):
75
+ os.makedirs(path)
76
+
77
+
78
+ def mkdirs(paths):
79
+ if isinstance(paths, str):
80
+ mkdir(paths)
81
+ else:
82
+ for path in paths:
83
+ mkdir(path)
84
+
85
+
86
+ def mkdir_and_rename(path):
87
+ if os.path.exists(path):
88
+ new_name = path + '_archived_' + get_timestamp()
89
+ print('Path already exists. Rename it to [{:s}]'.format(new_name))
90
+ os.rename(path, new_name)
91
+ os.makedirs(path)
92
+
93
+
94
+ '''
95
+ # =======================================
96
+ # read image from path
97
+ # Note: opencv is fast
98
+ # but read BGR numpy image
99
+ # =======================================
100
+ '''
101
+
102
+
103
+ # ----------------------------------------
104
+ # get single image of size HxWxn_channles (BGR)
105
+ # ----------------------------------------
106
+ def read_img(path):
107
+ # read image by cv2
108
+ # return: Numpy float32, HWC, BGR, [0,1]
109
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # cv2.IMREAD_GRAYSCALE
110
+ img = img.astype(np.float32) / 255.
111
+ if img.ndim == 2:
112
+ img = np.expand_dims(img, axis=2)
113
+ # some images have 4 channels
114
+ if img.shape[2] > 3:
115
+ img = img[:, :, :3]
116
+ return img
117
+
118
+
119
+ # ----------------------------------------
120
+ # get uint8 image of size HxWxn_channles (RGB)
121
+ # ----------------------------------------
122
+ def imread_uint(path, n_channels=3):
123
+ # input: path
124
+ # output: HxWx3(RGB or GGG), or HxWx1 (G)
125
+ if n_channels == 1:
126
+ img = cv2.imread(path, 0) # cv2.IMREAD_GRAYSCALE
127
+ img = np.expand_dims(img, axis=2) # HxWx1
128
+ elif n_channels == 3:
129
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # BGR or G
130
+ if img.ndim == 2:
131
+ img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # GGG
132
+ else:
133
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # RGB
134
+ return img
135
+
136
+
137
+ def imsave(img, img_path):
138
+ img = np.squeeze(img)
139
+ if img.ndim == 3:
140
+ img = img[:, :, [2, 1, 0]]
141
+ cv2.imwrite(img_path, img)
142
+
143
+
144
+ '''
145
+ # =======================================
146
+ # numpy(single) <---> numpy(uint)
147
+ # numpy(single) <---> tensor
148
+ # numpy(uint) <---> tensor
149
+ # =======================================
150
+ '''
151
+
152
+
153
+ # --------------------------------
154
+ # numpy(single) <---> numpy(uint)
155
+ # --------------------------------
156
+
157
+
158
+ def uint2single(img):
159
+
160
+ return np.float32(img/255.)
161
+
162
+
163
+ def uint2single1(img):
164
+
165
+ return np.float32(np.squeeze(img)/255.)
166
+
167
+
168
+ def single2uint(img):
169
+
170
+ return np.uint8((img.clip(0, 1)*255.).round())
171
+
172
+
173
+ def uint162single(img):
174
+
175
+ return np.float32(img/65535.)
176
+
177
+
178
+ def single2uint16(img):
179
+
180
+ return np.uint8((img.clip(0, 1)*65535.).round())
181
+
182
+
183
+ # --------------------------------
184
+ # numpy(uint) <---> tensor
185
+ # uint (HxWxn_channels (RGB) or G)
186
+ # --------------------------------
187
+
188
+
189
+ # convert uint (HxWxn_channels) to 4-dimensional torch tensor
190
+ def uint2tensor4(img):
191
+ if img.ndim == 2:
192
+ img = np.expand_dims(img, axis=2)
193
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.).unsqueeze(0)
194
+
195
+
196
+ # convert uint (HxWxn_channels) to 3-dimensional torch tensor
197
+ def uint2tensor3(img):
198
+ if img.ndim == 2:
199
+ img = np.expand_dims(img, axis=2)
200
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.)
201
+
202
+
203
+ # convert torch tensor to uint
204
+ def tensor2uint(img):
205
+ img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
206
+ if img.ndim == 3:
207
+ img = np.transpose(img, (1, 2, 0))
208
+ return np.uint8((img*255.0).round())
209
+
210
+
211
+ # --------------------------------
212
+ # numpy(single) <---> tensor
213
+ # single (HxWxn_channels (RGB) or G)
214
+ # --------------------------------
215
+
216
+
217
+ # convert single (HxWxn_channels) to 4-dimensional torch tensor
218
+ def single2tensor4(img):
219
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().unsqueeze(0)
220
+
221
+
222
+ # convert single (HxWxn_channels) to 3-dimensional torch tensor
223
+ def single2tensor3(img):
224
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float()
225
+
226
+
227
+ # convert torch tensor to single
228
+ def tensor2single(img):
229
+ img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
230
+ if img.ndim == 3:
231
+ img = np.transpose(img, (1, 2, 0))
232
+
233
+ return img
234
+
235
+ def tensor2single3(img):
236
+ img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
237
+ if img.ndim == 3:
238
+ img = np.transpose(img, (1, 2, 0))
239
+ elif img.ndim == 2:
240
+ img = np.expand_dims(img, axis=2)
241
+ return img
242
+
243
+
244
+ # from skimage.io import imread, imsave
245
+ def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
246
+ '''
247
+ Converts a torch Tensor into an image Numpy array of BGR channel order
248
+ Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
249
+ Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
250
+ '''
251
+ tensor = tensor.squeeze().float().cpu().clamp_(*min_max) # squeeze first, then clamp
252
+ tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) # to range [0,1]
253
+ n_dim = tensor.dim()
254
+ if n_dim == 4:
255
+ n_img = len(tensor)
256
+ img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()
257
+ img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
258
+ elif n_dim == 3:
259
+ img_np = tensor.numpy()
260
+ img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
261
+ elif n_dim == 2:
262
+ img_np = tensor.numpy()
263
+ else:
264
+ raise TypeError(
265
+ 'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
266
+ if out_type == np.uint8:
267
+ img_np = (img_np * 255.0).round()
268
+ # Important. Unlike matlab, numpy.uint8() WILL NOT round by default.
269
+ return img_np.astype(out_type)
270
+
271
+
272
+ '''
273
+ # =======================================
274
+ # image processing process on numpy image
275
+ # augment(img_list, hflip=True, rot=True):
276
+ # =======================================
277
+ '''
278
+
279
+
280
+ def augment_img(img, mode=0):
281
+ if mode == 0:
282
+ return img
283
+ elif mode == 1:
284
+ return np.flipud(np.rot90(img))
285
+ elif mode == 2:
286
+ return np.flipud(img)
287
+ elif mode == 3:
288
+ return np.rot90(img, k=3)
289
+ elif mode == 4:
290
+ return np.flipud(np.rot90(img, k=2))
291
+ elif mode == 5:
292
+ return np.rot90(img)
293
+ elif mode == 6:
294
+ return np.rot90(img, k=2)
295
+ elif mode == 7:
296
+ return np.flipud(np.rot90(img, k=3))
297
+
298
+
299
+ def augment_img_np3(img, mode=0):
300
+ if mode == 0:
301
+ return img
302
+ elif mode == 1:
303
+ return img.transpose(1, 0, 2)
304
+ elif mode == 2:
305
+ return img[::-1, :, :]
306
+ elif mode == 3:
307
+ img = img[::-1, :, :]
308
+ img = img.transpose(1, 0, 2)
309
+ return img
310
+ elif mode == 4:
311
+ return img[:, ::-1, :]
312
+ elif mode == 5:
313
+ img = img[:, ::-1, :]
314
+ img = img.transpose(1, 0, 2)
315
+ return img
316
+ elif mode == 6:
317
+ img = img[:, ::-1, :]
318
+ img = img[::-1, :, :]
319
+ return img
320
+ elif mode == 7:
321
+ img = img[:, ::-1, :]
322
+ img = img[::-1, :, :]
323
+ img = img.transpose(1, 0, 2)
324
+ return img
325
+
326
+
327
+ def augment_img_tensor(img, mode=0):
328
+ img_size = img.size()
329
+ img_np = img.data.cpu().numpy()
330
+ if len(img_size) == 3:
331
+ img_np = np.transpose(img_np, (1, 2, 0))
332
+ elif len(img_size) == 4:
333
+ img_np = np.transpose(img_np, (2, 3, 1, 0))
334
+ img_np = augment_img(img_np, mode=mode)
335
+ img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))
336
+ if len(img_size) == 3:
337
+ img_tensor = img_tensor.permute(2, 0, 1)
338
+ elif len(img_size) == 4:
339
+ img_tensor = img_tensor.permute(3, 2, 0, 1)
340
+
341
+ return img_tensor.type_as(img)
342
+
343
+
344
+ def augment_imgs(img_list, hflip=True, rot=True):
345
+ # horizontal flip OR rotate
346
+ hflip = hflip and random.random() < 0.5
347
+ vflip = rot and random.random() < 0.5
348
+ rot90 = rot and random.random() < 0.5
349
+
350
+ def _augment(img):
351
+ if hflip:
352
+ img = img[:, ::-1, :]
353
+ if vflip:
354
+ img = img[::-1, :, :]
355
+ if rot90:
356
+ img = img.transpose(1, 0, 2)
357
+ return img
358
+
359
+ return [_augment(img) for img in img_list]
360
+
361
+
362
+ '''
363
+ # =======================================
364
+ # image processing process on numpy image
365
+ # channel_convert(in_c, tar_type, img_list):
366
+ # rgb2ycbcr(img, only_y=True):
367
+ # bgr2ycbcr(img, only_y=True):
368
+ # ycbcr2rgb(img):
369
+ # modcrop(img_in, scale):
370
+ # =======================================
371
+ '''
372
+
373
+
374
+ def rgb2ycbcr(img, only_y=True):
375
+ '''same as matlab rgb2ycbcr
376
+ only_y: only return Y channel
377
+ Input:
378
+ uint8, [0, 255]
379
+ float, [0, 1]
380
+ '''
381
+ in_img_type = img.dtype
382
+ img.astype(np.float32)
383
+ if in_img_type != np.uint8:
384
+ img *= 255.
385
+ # convert
386
+ if only_y:
387
+ rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
388
+ else:
389
+ rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
390
+ [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
391
+ if in_img_type == np.uint8:
392
+ rlt = rlt.round()
393
+ else:
394
+ rlt /= 255.
395
+ return rlt.astype(in_img_type)
396
+
397
+
398
+ def ycbcr2rgb(img):
399
+ '''same as matlab ycbcr2rgb
400
+ Input:
401
+ uint8, [0, 255]
402
+ float, [0, 1]
403
+ '''
404
+ in_img_type = img.dtype
405
+ img.astype(np.float32)
406
+ if in_img_type != np.uint8:
407
+ img *= 255.
408
+ # convert
409
+ rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
410
+ [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
411
+ if in_img_type == np.uint8:
412
+ rlt = rlt.round()
413
+ else:
414
+ rlt /= 255.
415
+ return rlt.astype(in_img_type)
416
+
417
+
418
+ def bgr2ycbcr(img, only_y=True):
419
+ '''bgr version of rgb2ycbcr
420
+ only_y: only return Y channel
421
+ Input:
422
+ uint8, [0, 255]
423
+ float, [0, 1]
424
+ '''
425
+ in_img_type = img.dtype
426
+ img.astype(np.float32)
427
+ if in_img_type != np.uint8:
428
+ img *= 255.
429
+ # convert
430
+ if only_y:
431
+ rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
432
+ else:
433
+ rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
434
+ [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
435
+ if in_img_type == np.uint8:
436
+ rlt = rlt.round()
437
+ else:
438
+ rlt /= 255.
439
+ return rlt.astype(in_img_type)
440
+
441
+
442
+ def modcrop(img_in, scale):
443
+ # img_in: Numpy, HWC or HW
444
+ img = np.copy(img_in)
445
+ if img.ndim == 2:
446
+ H, W = img.shape
447
+ H_r, W_r = H % scale, W % scale
448
+ img = img[:H - H_r, :W - W_r]
449
+ elif img.ndim == 3:
450
+ H, W, C = img.shape
451
+ H_r, W_r = H % scale, W % scale
452
+ img = img[:H - H_r, :W - W_r, :]
453
+ else:
454
+ raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
455
+ return img
456
+
457
+
458
+ def shave(img_in, border=0):
459
+ # img_in: Numpy, HWC or HW
460
+ img = np.copy(img_in)
461
+ h, w = img.shape[:2]
462
+ img = img[border:h-border, border:w-border]
463
+ return img
464
+
465
+
466
+ def channel_convert(in_c, tar_type, img_list):
467
+ # conversion among BGR, gray and y
468
+ if in_c == 3 and tar_type == 'gray': # BGR to gray
469
+ gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
470
+ return [np.expand_dims(img, axis=2) for img in gray_list]
471
+ elif in_c == 3 and tar_type == 'y': # BGR to y
472
+ y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
473
+ return [np.expand_dims(img, axis=2) for img in y_list]
474
+ elif in_c == 1 and tar_type == 'RGB': # gray/y to BGR
475
+ return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
476
+ else:
477
+ return img_list
478
+
479
+
480
+ '''
481
+ # =======================================
482
+ # metric, PSNR and SSIM
483
+ # =======================================
484
+ '''
485
+
486
+
487
+ # ----------
488
+ # PSNR
489
+ # ----------
490
+ def calculate_psnr(img1, img2, border=0):
491
+ # img1 and img2 have range [0, 255]
492
+ if not img1.shape == img2.shape:
493
+ raise ValueError('Input images must have the same dimensions.')
494
+ h, w = img1.shape[:2]
495
+ img1 = img1[border:h-border, border:w-border]
496
+ img2 = img2[border:h-border, border:w-border]
497
+
498
+ img1 = img1.astype(np.float64)
499
+ img2 = img2.astype(np.float64)
500
+ mse = np.mean((img1 - img2)**2)
501
+ if mse == 0:
502
+ return float('inf')
503
+ return 20 * math.log10(255.0 / math.sqrt(mse))
504
+
505
+
506
+ # ----------
507
+ # SSIM
508
+ # ----------
509
+ def calculate_ssim(img1, img2, border=0):
510
+ '''calculate SSIM
511
+ the same outputs as MATLAB's
512
+ img1, img2: [0, 255]
513
+ '''
514
+ if not img1.shape == img2.shape:
515
+ raise ValueError('Input images must have the same dimensions.')
516
+ h, w = img1.shape[:2]
517
+ img1 = img1[border:h-border, border:w-border]
518
+ img2 = img2[border:h-border, border:w-border]
519
+
520
+ if img1.ndim == 2:
521
+ return ssim(img1, img2)
522
+ elif img1.ndim == 3:
523
+ if img1.shape[2] == 3:
524
+ ssims = []
525
+ for i in range(3):
526
+ ssims.append(ssim(img1, img2))
527
+ return np.array(ssims).mean()
528
+ elif img1.shape[2] == 1:
529
+ return ssim(np.squeeze(img1), np.squeeze(img2))
530
+ else:
531
+ raise ValueError('Wrong input image dimensions.')
532
+
533
+
534
+ def ssim(img1, img2):
535
+ C1 = (0.01 * 255)**2
536
+ C2 = (0.03 * 255)**2
537
+
538
+ img1 = img1.astype(np.float64)
539
+ img2 = img2.astype(np.float64)
540
+ kernel = cv2.getGaussianKernel(11, 1.5)
541
+ window = np.outer(kernel, kernel.transpose())
542
+
543
+ mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid
544
+ mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
545
+ mu1_sq = mu1**2
546
+ mu2_sq = mu2**2
547
+ mu1_mu2 = mu1 * mu2
548
+ sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
549
+ sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
550
+ sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
551
+
552
+ ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
553
+ (sigma1_sq + sigma2_sq + C2))
554
+ return ssim_map.mean()
555
+
556
+
557
+ '''
558
+ # =======================================
559
+ # pytorch version of matlab imresize
560
+ # =======================================
561
+ '''
562
+
563
+
564
+ # matlab 'imresize' function, now only support 'bicubic'
565
+ def cubic(x):
566
+ absx = torch.abs(x)
567
+ absx2 = absx**2
568
+ absx3 = absx**3
569
+ return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \
570
+ (-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx))
571
+
572
+
573
+ def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
574
+ if (scale < 1) and (antialiasing):
575
+ # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
576
+ kernel_width = kernel_width / scale
577
+
578
+ # Output-space coordinates
579
+ x = torch.linspace(1, out_length, out_length)
580
+
581
+ # Input-space coordinates. Calculate the inverse mapping such that 0.5
582
+ # in output space maps to 0.5 in input space, and 0.5+scale in output
583
+ # space maps to 1.5 in input space.
584
+ u = x / scale + 0.5 * (1 - 1 / scale)
585
+
586
+ # What is the left-most pixel that can be involved in the computation?
587
+ left = torch.floor(u - kernel_width / 2)
588
+
589
+ # What is the maximum number of pixels that can be involved in the
590
+ # computation? Note: it's OK to use an extra pixel here; if the
591
+ # corresponding weights are all zero, it will be eliminated at the end
592
+ # of this function.
593
+ P = math.ceil(kernel_width) + 2
594
+
595
+ # The indices of the input pixels involved in computing the k-th output
596
+ # pixel are in row k of the indices matrix.
597
+ indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
598
+ 1, P).expand(out_length, P)
599
+
600
+ # The weights used to compute the k-th output pixel are in row k of the
601
+ # weights matrix.
602
+ distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
603
+ # apply cubic kernel
604
+ if (scale < 1) and (antialiasing):
605
+ weights = scale * cubic(distance_to_center * scale)
606
+ else:
607
+ weights = cubic(distance_to_center)
608
+ # Normalize the weights matrix so that each row sums to 1.
609
+ weights_sum = torch.sum(weights, 1).view(out_length, 1)
610
+ weights = weights / weights_sum.expand(out_length, P)
611
+
612
+ # If a column in weights is all zero, get rid of it. only consider the first and last column.
613
+ weights_zero_tmp = torch.sum((weights == 0), 0)
614
+ if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
615
+ indices = indices.narrow(1, 1, P - 2)
616
+ weights = weights.narrow(1, 1, P - 2)
617
+ if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
618
+ indices = indices.narrow(1, 0, P - 2)
619
+ weights = weights.narrow(1, 0, P - 2)
620
+ weights = weights.contiguous()
621
+ indices = indices.contiguous()
622
+ sym_len_s = -indices.min() + 1
623
+ sym_len_e = indices.max() - in_length
624
+ indices = indices + sym_len_s - 1
625
+ return weights, indices, int(sym_len_s), int(sym_len_e)
626
+
627
+
628
+ # --------------------------------
629
+ # imresize for tensor image
630
+ # --------------------------------
631
+ def imresize(img, scale, antialiasing=True):
632
+ # Now the scale should be the same for H and W
633
+ # input: img: pytorch tensor, CHW or HW [0,1]
634
+ # output: CHW or HW [0,1] w/o round
635
+ need_squeeze = True if img.dim() == 2 else False
636
+ if need_squeeze:
637
+ img.unsqueeze_(0)
638
+ in_C, in_H, in_W = img.size()
639
+ out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
640
+ kernel_width = 4
641
+ kernel = 'cubic'
642
+
643
+ # Return the desired dimension order for performing the resize. The
644
+ # strategy is to perform the resize first along the dimension with the
645
+ # smallest scale factor.
646
+ # Now we do not support this.
647
+
648
+ # get weights and indices
649
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
650
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
651
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
652
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
653
+ # process H dimension
654
+ # symmetric copying
655
+ img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
656
+ img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
657
+
658
+ sym_patch = img[:, :sym_len_Hs, :]
659
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
660
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
661
+ img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
662
+
663
+ sym_patch = img[:, -sym_len_He:, :]
664
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
665
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
666
+ img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
667
+
668
+ out_1 = torch.FloatTensor(in_C, out_H, in_W)
669
+ kernel_width = weights_H.size(1)
670
+ for i in range(out_H):
671
+ idx = int(indices_H[i][0])
672
+ for j in range(out_C):
673
+ out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
674
+
675
+ # process W dimension
676
+ # symmetric copying
677
+ out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
678
+ out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
679
+
680
+ sym_patch = out_1[:, :, :sym_len_Ws]
681
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
682
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
683
+ out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
684
+
685
+ sym_patch = out_1[:, :, -sym_len_We:]
686
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
687
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
688
+ out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
689
+
690
+ out_2 = torch.FloatTensor(in_C, out_H, out_W)
691
+ kernel_width = weights_W.size(1)
692
+ for i in range(out_W):
693
+ idx = int(indices_W[i][0])
694
+ for j in range(out_C):
695
+ out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_W[i])
696
+ if need_squeeze:
697
+ out_2.squeeze_()
698
+ return out_2
699
+
700
+
701
+ # --------------------------------
702
+ # imresize for numpy image
703
+ # --------------------------------
704
+ def imresize_np(img, scale, antialiasing=True):
705
+ # Now the scale should be the same for H and W
706
+ # input: img: Numpy, HWC or HW [0,1]
707
+ # output: HWC or HW [0,1] w/o round
708
+ img = torch.from_numpy(img)
709
+ need_squeeze = True if img.dim() == 2 else False
710
+ if need_squeeze:
711
+ img.unsqueeze_(2)
712
+
713
+ in_H, in_W, in_C = img.size()
714
+ out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
715
+ kernel_width = 4
716
+ kernel = 'cubic'
717
+
718
+ # Return the desired dimension order for performing the resize. The
719
+ # strategy is to perform the resize first along the dimension with the
720
+ # smallest scale factor.
721
+ # Now we do not support this.
722
+
723
+ # get weights and indices
724
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
725
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
726
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
727
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
728
+ # process H dimension
729
+ # symmetric copying
730
+ img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
731
+ img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
732
+
733
+ sym_patch = img[:sym_len_Hs, :, :]
734
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
735
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
736
+ img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
737
+
738
+ sym_patch = img[-sym_len_He:, :, :]
739
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
740
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
741
+ img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
742
+
743
+ out_1 = torch.FloatTensor(out_H, in_W, in_C)
744
+ kernel_width = weights_H.size(1)
745
+ for i in range(out_H):
746
+ idx = int(indices_H[i][0])
747
+ for j in range(out_C):
748
+ out_1[i, :, j] = img_aug[idx:idx + kernel_width, :, j].transpose(0, 1).mv(weights_H[i])
749
+
750
+ # process W dimension
751
+ # symmetric copying
752
+ out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
753
+ out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
754
+
755
+ sym_patch = out_1[:, :sym_len_Ws, :]
756
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
757
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
758
+ out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
759
+
760
+ sym_patch = out_1[:, -sym_len_We:, :]
761
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
762
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
763
+ out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
764
+
765
+ out_2 = torch.FloatTensor(out_H, out_W, in_C)
766
+ kernel_width = weights_W.size(1)
767
+ for i in range(out_W):
768
+ idx = int(indices_W[i][0])
769
+ for j in range(out_C):
770
+ out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width, j].mv(weights_W[i])
771
+ if need_squeeze:
772
+ out_2.squeeze_()
773
+
774
+ return out_2.numpy()
775
+
776
+
777
+ if __name__ == '__main__':
778
+ img = imread_uint('test.bmp',3)
utils/utils_logger.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ import datetime
4
+ import logging
5
+
6
+
7
+ '''
8
+ modified by Kai Zhang (github: https://github.com/cszn)
9
+ 03/03/2019
10
+ https://github.com/xinntao/BasicSR
11
+ '''
12
+
13
+
14
+ def log(*args, **kwargs):
15
+ print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), *args, **kwargs)
16
+
17
+
18
+ '''
19
+ # ===============================
20
+ # logger
21
+ # logger_name = None = 'base' ???
22
+ # ===============================
23
+ '''
24
+
25
+
26
+ def logger_info(logger_name, log_path='default_logger.log'):
27
+ ''' set up logger
28
+ modified by Kai Zhang (github: https://github.com/cszn)
29
+ '''
30
+ log = logging.getLogger(logger_name)
31
+ if log.hasHandlers():
32
+ print('LogHandlers exist!')
33
+ else:
34
+ print('LogHandlers setup!')
35
+ level = logging.INFO
36
+ formatter = logging.Formatter('%(asctime)s.%(msecs)03d : %(message)s', datefmt='%y-%m-%d %H:%M:%S')
37
+ fh = logging.FileHandler(log_path, mode='a')
38
+ fh.setFormatter(formatter)
39
+ log.setLevel(level)
40
+ log.addHandler(fh)
41
+ # print(len(log.handlers))
42
+
43
+ sh = logging.StreamHandler()
44
+ sh.setFormatter(formatter)
45
+ log.addHandler(sh)
46
+
47
+
48
+ '''
49
+ # ===============================
50
+ # print to file and std_out simultaneously
51
+ # ===============================
52
+ '''
53
+
54
+
55
+ class logger_print(object):
56
+ def __init__(self, log_path="default.log"):
57
+ self.terminal = sys.stdout
58
+ self.log = open(log_path, 'a')
59
+
60
+ def write(self, message):
61
+ self.terminal.write(message)
62
+ self.log.write(message) # write the message
63
+
64
+ def flush(self):
65
+ pass