doevent commited on
Commit
25ef88c
1 Parent(s): f3f6927

Upload utils/util_calculate_psnr_ssim.py

Browse files
Files changed (1) hide show
  1. utils/util_calculate_psnr_ssim.py +320 -0
utils/util_calculate_psnr_ssim.py ADDED
@@ -0,0 +1,320 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -----------------------------------------------------------------------------------
2
+ # https://github.com/JingyunLiang/SwinIR/blob/main/utils/util_calculate_psnr_ssim.py
3
+ # -----------------------------------------------------------------------------------
4
+
5
+ import cv2
6
+ import torch
7
+ import numpy as np
8
+
9
+ def calculate_psnr(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
10
+ """Calculate PSNR (Peak Signal-to-Noise Ratio).
11
+ Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
12
+ Args:
13
+ img1 (ndarray): Images with range [0, 255].
14
+ img2 (ndarray): Images with range [0, 255].
15
+ crop_border (int): Cropped pixels in each edge of an image. These
16
+ pixels are not involved in the PSNR calculation.
17
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
18
+ Default: 'HWC'.
19
+ test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
20
+ Returns:
21
+ float: psnr result.
22
+ """
23
+
24
+ assert img1.shape == img2.shape, (f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
25
+ if input_order not in ['HWC', 'CHW']:
26
+ raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
27
+ img1 = reorder_image(img1, input_order=input_order)
28
+ img2 = reorder_image(img2, input_order=input_order)
29
+ img1 = img1.astype(np.float64)
30
+ img2 = img2.astype(np.float64)
31
+
32
+ if crop_border != 0:
33
+ img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
34
+ img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
35
+
36
+ if test_y_channel:
37
+ img1 = to_y_channel(img1)
38
+ img2 = to_y_channel(img2)
39
+
40
+ mse = np.mean((img1 - img2) ** 2)
41
+ if mse == 0:
42
+ return float('inf')
43
+ return 20. * np.log10(255. / np.sqrt(mse))
44
+
45
+
46
+ def _ssim(img1, img2):
47
+ """Calculate SSIM (structural similarity) for one channel images.
48
+ It is called by func:`calculate_ssim`.
49
+ Args:
50
+ img1 (ndarray): Images with range [0, 255] with order 'HWC'.
51
+ img2 (ndarray): Images with range [0, 255] with order 'HWC'.
52
+ Returns:
53
+ float: ssim result.
54
+ """
55
+
56
+ C1 = (0.01 * 255) ** 2
57
+ C2 = (0.03 * 255) ** 2
58
+
59
+ img1 = img1.astype(np.float64)
60
+ img2 = img2.astype(np.float64)
61
+ kernel = cv2.getGaussianKernel(11, 1.5)
62
+ window = np.outer(kernel, kernel.transpose())
63
+
64
+ mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
65
+ mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
66
+ mu1_sq = mu1 ** 2
67
+ mu2_sq = mu2 ** 2
68
+ mu1_mu2 = mu1 * mu2
69
+ sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
70
+ sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
71
+ sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
72
+
73
+ ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
74
+ return ssim_map.mean()
75
+
76
+
77
+ def calculate_ssim(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
78
+ """Calculate SSIM (structural similarity).
79
+ Ref:
80
+ Image quality assessment: From error visibility to structural similarity
81
+ The results are the same as that of the official released MATLAB code in
82
+ https://ece.uwaterloo.ca/~z70wang/research/ssim/.
83
+ For three-channel images, SSIM is calculated for each channel and then
84
+ averaged.
85
+ Args:
86
+ img1 (ndarray): Images with range [0, 255].
87
+ img2 (ndarray): Images with range [0, 255].
88
+ crop_border (int): Cropped pixels in each edge of an image. These
89
+ pixels are not involved in the SSIM calculation.
90
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
91
+ Default: 'HWC'.
92
+ test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
93
+ Returns:
94
+ float: ssim result.
95
+ """
96
+
97
+ assert img1.shape == img2.shape, (f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
98
+ if input_order not in ['HWC', 'CHW']:
99
+ raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
100
+ img1 = reorder_image(img1, input_order=input_order)
101
+ img2 = reorder_image(img2, input_order=input_order)
102
+ img1 = img1.astype(np.float64)
103
+ img2 = img2.astype(np.float64)
104
+
105
+ if crop_border != 0:
106
+ img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
107
+ img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
108
+
109
+ if test_y_channel:
110
+ img1 = to_y_channel(img1)
111
+ img2 = to_y_channel(img2)
112
+
113
+ ssims = []
114
+ for i in range(img1.shape[2]):
115
+ ssims.append(_ssim(img1[..., i], img2[..., i]))
116
+ return np.array(ssims).mean()
117
+
118
+
119
+ def _blocking_effect_factor(im):
120
+ block_size = 8
121
+
122
+ block_horizontal_positions = torch.arange(7, im.shape[3] - 1, 8)
123
+ block_vertical_positions = torch.arange(7, im.shape[2] - 1, 8)
124
+
125
+ horizontal_block_difference = (
126
+ (im[:, :, :, block_horizontal_positions] - im[:, :, :, block_horizontal_positions + 1]) ** 2).sum(
127
+ 3).sum(2).sum(1)
128
+ vertical_block_difference = (
129
+ (im[:, :, block_vertical_positions, :] - im[:, :, block_vertical_positions + 1, :]) ** 2).sum(3).sum(
130
+ 2).sum(1)
131
+
132
+ nonblock_horizontal_positions = np.setdiff1d(torch.arange(0, im.shape[3] - 1), block_horizontal_positions)
133
+ nonblock_vertical_positions = np.setdiff1d(torch.arange(0, im.shape[2] - 1), block_vertical_positions)
134
+
135
+ horizontal_nonblock_difference = (
136
+ (im[:, :, :, nonblock_horizontal_positions] - im[:, :, :, nonblock_horizontal_positions + 1]) ** 2).sum(
137
+ 3).sum(2).sum(1)
138
+ vertical_nonblock_difference = (
139
+ (im[:, :, nonblock_vertical_positions, :] - im[:, :, nonblock_vertical_positions + 1, :]) ** 2).sum(
140
+ 3).sum(2).sum(1)
141
+
142
+ n_boundary_horiz = im.shape[2] * (im.shape[3] // block_size - 1)
143
+ n_boundary_vert = im.shape[3] * (im.shape[2] // block_size - 1)
144
+ boundary_difference = (horizontal_block_difference + vertical_block_difference) / (
145
+ n_boundary_horiz + n_boundary_vert)
146
+
147
+ n_nonboundary_horiz = im.shape[2] * (im.shape[3] - 1) - n_boundary_horiz
148
+ n_nonboundary_vert = im.shape[3] * (im.shape[2] - 1) - n_boundary_vert
149
+ nonboundary_difference = (horizontal_nonblock_difference + vertical_nonblock_difference) / (
150
+ n_nonboundary_horiz + n_nonboundary_vert)
151
+
152
+ scaler = np.log2(block_size) / np.log2(min([im.shape[2], im.shape[3]]))
153
+ bef = scaler * (boundary_difference - nonboundary_difference)
154
+
155
+ bef[boundary_difference <= nonboundary_difference] = 0
156
+ return bef
157
+
158
+
159
+ def calculate_psnrb(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
160
+ """Calculate PSNR-B (Peak Signal-to-Noise Ratio).
161
+ Ref: Quality assessment of deblocked images, for JPEG image deblocking evaluation
162
+ # https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py
163
+ Args:
164
+ img1 (ndarray): Images with range [0, 255].
165
+ img2 (ndarray): Images with range [0, 255].
166
+ crop_border (int): Cropped pixels in each edge of an image. These
167
+ pixels are not involved in the PSNR calculation.
168
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
169
+ Default: 'HWC'.
170
+ test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
171
+ Returns:
172
+ float: psnr result.
173
+ """
174
+
175
+ assert img1.shape == img2.shape, (f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
176
+ if input_order not in ['HWC', 'CHW']:
177
+ raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
178
+ img1 = reorder_image(img1, input_order=input_order)
179
+ img2 = reorder_image(img2, input_order=input_order)
180
+ img1 = img1.astype(np.float64)
181
+ img2 = img2.astype(np.float64)
182
+
183
+ if crop_border != 0:
184
+ img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
185
+ img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
186
+
187
+ if test_y_channel:
188
+ img1 = to_y_channel(img1)
189
+ img2 = to_y_channel(img2)
190
+
191
+ # follow https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py
192
+ img1 = torch.from_numpy(img1).permute(2, 0, 1).unsqueeze(0) / 255.
193
+ img2 = torch.from_numpy(img2).permute(2, 0, 1).unsqueeze(0) / 255.
194
+
195
+ total = 0
196
+ for c in range(img1.shape[1]):
197
+ mse = torch.nn.functional.mse_loss(img1[:, c:c + 1, :, :], img2[:, c:c + 1, :, :], reduction='none')
198
+ bef = _blocking_effect_factor(img1[:, c:c + 1, :, :])
199
+
200
+ mse = mse.view(mse.shape[0], -1).mean(1)
201
+ total += 10 * torch.log10(1 / (mse + bef))
202
+
203
+ return float(total) / img1.shape[1]
204
+
205
+
206
+ def reorder_image(img, input_order='HWC'):
207
+ """Reorder images to 'HWC' order.
208
+ If the input_order is (h, w), return (h, w, 1);
209
+ If the input_order is (c, h, w), return (h, w, c);
210
+ If the input_order is (h, w, c), return as it is.
211
+ Args:
212
+ img (ndarray): Input image.
213
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
214
+ If the input image shape is (h, w), input_order will not have
215
+ effects. Default: 'HWC'.
216
+ Returns:
217
+ ndarray: reordered image.
218
+ """
219
+
220
+ if input_order not in ['HWC', 'CHW']:
221
+ raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' "'HWC' and 'CHW'")
222
+ if len(img.shape) == 2:
223
+ img = img[..., None]
224
+ if input_order == 'CHW':
225
+ img = img.transpose(1, 2, 0)
226
+ return img
227
+
228
+
229
+ def to_y_channel(img):
230
+ """Change to Y channel of YCbCr.
231
+ Args:
232
+ img (ndarray): Images with range [0, 255].
233
+ Returns:
234
+ (ndarray): Images with range [0, 255] (float type) without round.
235
+ """
236
+ img = img.astype(np.float32) / 255.
237
+ if img.ndim == 3 and img.shape[2] == 3:
238
+ img = bgr2ycbcr(img, y_only=True)
239
+ img = img[..., None]
240
+ return img * 255.
241
+
242
+
243
+ def _convert_input_type_range(img):
244
+ """Convert the type and range of the input image.
245
+ It converts the input image to np.float32 type and range of [0, 1].
246
+ It is mainly used for pre-processing the input image in colorspace
247
+ convertion functions such as rgb2ycbcr and ycbcr2rgb.
248
+ Args:
249
+ img (ndarray): The input image. It accepts:
250
+ 1. np.uint8 type with range [0, 255];
251
+ 2. np.float32 type with range [0, 1].
252
+ Returns:
253
+ (ndarray): The converted image with type of np.float32 and range of
254
+ [0, 1].
255
+ """
256
+ img_type = img.dtype
257
+ img = img.astype(np.float32)
258
+ if img_type == np.float32:
259
+ pass
260
+ elif img_type == np.uint8:
261
+ img /= 255.
262
+ else:
263
+ raise TypeError('The img type should be np.float32 or np.uint8, ' f'but got {img_type}')
264
+ return img
265
+
266
+
267
+ def _convert_output_type_range(img, dst_type):
268
+ """Convert the type and range of the image according to dst_type.
269
+ It converts the image to desired type and range. If `dst_type` is np.uint8,
270
+ images will be converted to np.uint8 type with range [0, 255]. If
271
+ `dst_type` is np.float32, it converts the image to np.float32 type with
272
+ range [0, 1].
273
+ It is mainly used for post-processing images in colorspace convertion
274
+ functions such as rgb2ycbcr and ycbcr2rgb.
275
+ Args:
276
+ img (ndarray): The image to be converted with np.float32 type and
277
+ range [0, 255].
278
+ dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
279
+ converts the image to np.uint8 type with range [0, 255]. If
280
+ dst_type is np.float32, it converts the image to np.float32 type
281
+ with range [0, 1].
282
+ Returns:
283
+ (ndarray): The converted image with desired type and range.
284
+ """
285
+ if dst_type not in (np.uint8, np.float32):
286
+ raise TypeError('The dst_type should be np.float32 or np.uint8, ' f'but got {dst_type}')
287
+ if dst_type == np.uint8:
288
+ img = img.round()
289
+ else:
290
+ img /= 255.
291
+ return img.astype(dst_type)
292
+
293
+
294
+ def bgr2ycbcr(img, y_only=False):
295
+ """Convert a BGR image to YCbCr image.
296
+ The bgr version of rgb2ycbcr.
297
+ It implements the ITU-R BT.601 conversion for standard-definition
298
+ television. See more details in
299
+ https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
300
+ It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
301
+ In OpenCV, it implements a JPEG conversion. See more details in
302
+ https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
303
+ Args:
304
+ img (ndarray): The input image. It accepts:
305
+ 1. np.uint8 type with range [0, 255];
306
+ 2. np.float32 type with range [0, 1].
307
+ y_only (bool): Whether to only return Y channel. Default: False.
308
+ Returns:
309
+ ndarray: The converted YCbCr image. The output image has the same type
310
+ and range as input image.
311
+ """
312
+ img_type = img.dtype
313
+ img = _convert_input_type_range(img)
314
+ if y_only:
315
+ out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
316
+ else:
317
+ out_img = np.matmul(
318
+ img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]
319
+ out_img = _convert_output_type_range(out_img, img_type)
320
+ return out_img