tomaseo2022 commited on
Commit
80e30a3
1 Parent(s): 518cfb2

Upload 6 files

Browse files
.gitattributes CHANGED
@@ -31,3 +31,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
31
  *.zip filter=lfs diff=lfs merge=lfs -text
32
  *.zst filter=lfs diff=lfs merge=lfs -text
33
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
31
  *.zip filter=lfs diff=lfs merge=lfs -text
32
  *.zst filter=lfs diff=lfs merge=lfs -text
33
  *tfevents* filter=lfs diff=lfs merge=lfs -text
34
+ ETH_LR.png filter=lfs diff=lfs merge=lfs -text
ETH_LR.png ADDED

Git LFS Details

  • SHA256: a6049b3f63961a85f28498a26fb1da726cd470d9d37d8fbc49c70ba197af3d63
  • Pointer size: 132 Bytes
  • Size of remote file: 1.67 MB
download-weights.sh ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/sh
2
+
3
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth -P experiments/pretrained_models
4
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/004_grayDN_DFWB_s128w8_SwinIR-M_noise15.pth -P experiments/pretrained_models
5
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/004_grayDN_DFWB_s128w8_SwinIR-M_noise25.pth -P experiments/pretrained_models
6
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/004_grayDN_DFWB_s128w8_SwinIR-M_noise50.pth -P experiments/pretrained_models
7
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/005_colorDN_DFWB_s128w8_SwinIR-M_noise15.pth -P experiments/pretrained_models
8
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth -P experiments/pretrained_models
9
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/005_colorDN_DFWB_s128w8_SwinIR-M_noise50.pth -P experiments/pretrained_models
10
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg10.pth -P experiments/pretrained_models
11
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg20.pth -P experiments/pretrained_models
12
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg30.pth -P experiments/pretrained_models
13
+ wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg40.pth -P experiments/pretrained_models
main_test_swinir.py ADDED
@@ -0,0 +1,253 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import cv2
3
+ import glob
4
+ import numpy as np
5
+ from collections import OrderedDict
6
+ import os
7
+ import torch
8
+ import requests
9
+
10
+ from models.network_swinir import SwinIR as net
11
+ from utils import util_calculate_psnr_ssim as util
12
+
13
+
14
+ def main():
15
+ parser = argparse.ArgumentParser()
16
+ parser.add_argument('--task', type=str, default='color_dn', help='classical_sr, lightweight_sr, real_sr, '
17
+ 'gray_dn, color_dn, jpeg_car')
18
+ parser.add_argument('--scale', type=int, default=1, help='scale factor: 1, 2, 3, 4, 8') # 1 for dn and jpeg car
19
+ parser.add_argument('--noise', type=int, default=15, help='noise level: 15, 25, 50')
20
+ parser.add_argument('--jpeg', type=int, default=40, help='scale factor: 10, 20, 30, 40')
21
+ parser.add_argument('--training_patch_size', type=int, default=128, help='patch size used in training SwinIR. '
22
+ 'Just used to differentiate two different settings in Table 2 of the paper. '
23
+ 'Images are NOT tested patch by patch.')
24
+ parser.add_argument('--large_model', action='store_true', help='use large model, only provided for real image sr')
25
+ parser.add_argument('--model_path', type=str,
26
+ default='model_zoo/swinir/001_classicalSR_DIV2K_s48w8_SwinIR-M_x2.pth')
27
+ parser.add_argument('--folder_lq', type=str, default=None, help='input low-quality test image folder')
28
+ parser.add_argument('--folder_gt', type=str, default=None, help='input ground-truth test image folder')
29
+ args = parser.parse_args()
30
+
31
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
32
+ # set up model
33
+ if os.path.exists(args.model_path):
34
+ print(f'loading model from {args.model_path}')
35
+ else:
36
+ os.makedirs(os.path.dirname(args.model_path), exist_ok=True)
37
+ url = 'https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/{}'.format(os.path.basename(args.model_path))
38
+ r = requests.get(url, allow_redirects=True)
39
+ print(f'downloading model {args.model_path}')
40
+ open(args.model_path, 'wb').write(r.content)
41
+
42
+ model = define_model(args)
43
+ model.eval()
44
+ model = model.to(device)
45
+
46
+ # setup folder and path
47
+ folder, save_dir, border, window_size = setup(args)
48
+ os.makedirs(save_dir, exist_ok=True)
49
+ test_results = OrderedDict()
50
+ test_results['psnr'] = []
51
+ test_results['ssim'] = []
52
+ test_results['psnr_y'] = []
53
+ test_results['ssim_y'] = []
54
+ test_results['psnr_b'] = []
55
+ psnr, ssim, psnr_y, ssim_y, psnr_b = 0, 0, 0, 0, 0
56
+
57
+ for idx, path in enumerate(sorted(glob.glob(os.path.join(folder, '*')))):
58
+ # read image
59
+ imgname, img_lq, img_gt = get_image_pair(args, path) # image to HWC-BGR, float32
60
+ img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1)) # HCW-BGR to CHW-RGB
61
+ img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device) # CHW-RGB to NCHW-RGB
62
+
63
+ # inference
64
+ with torch.no_grad():
65
+ # pad input image to be a multiple of window_size
66
+ _, _, h_old, w_old = img_lq.size()
67
+ h_pad = (h_old // window_size + 1) * window_size - h_old
68
+ w_pad = (w_old // window_size + 1) * window_size - w_old
69
+ img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
70
+ img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
71
+ output = model(img_lq)
72
+ output = output[..., :h_old * args.scale, :w_old * args.scale]
73
+
74
+ # save image
75
+ output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
76
+ if output.ndim == 3:
77
+ output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR
78
+ output = (output * 255.0).round().astype(np.uint8) # float32 to uint8
79
+ cv2.imwrite(f'{save_dir}/{imgname}_SwinIR.png', output)
80
+
81
+ # evaluate psnr/ssim/psnr_b
82
+ if img_gt is not None:
83
+ img_gt = (img_gt * 255.0).round().astype(np.uint8) # float32 to uint8
84
+ img_gt = img_gt[:h_old * args.scale, :w_old * args.scale, ...] # crop gt
85
+ img_gt = np.squeeze(img_gt)
86
+
87
+ psnr = util.calculate_psnr(output, img_gt, crop_border=border)
88
+ ssim = util.calculate_ssim(output, img_gt, crop_border=border)
89
+ test_results['psnr'].append(psnr)
90
+ test_results['ssim'].append(ssim)
91
+ if img_gt.ndim == 3: # RGB image
92
+ psnr_y = util.calculate_psnr(output, img_gt, crop_border=border, test_y_channel=True)
93
+ ssim_y = util.calculate_ssim(output, img_gt, crop_border=border, test_y_channel=True)
94
+ test_results['psnr_y'].append(psnr_y)
95
+ test_results['ssim_y'].append(ssim_y)
96
+ if args.task in ['jpeg_car']:
97
+ psnr_b = util.calculate_psnrb(output, img_gt, crop_border=border, test_y_channel=True)
98
+ test_results['psnr_b'].append(psnr_b)
99
+ print('Testing {:d} {:20s} - PSNR: {:.2f} dB; SSIM: {:.4f}; '
100
+ 'PSNR_Y: {:.2f} dB; SSIM_Y: {:.4f}; '
101
+ 'PSNR_B: {:.2f} dB.'.
102
+ format(idx, imgname, psnr, ssim, psnr_y, ssim_y, psnr_b))
103
+ else:
104
+ print('Testing {:d} {:20s}'.format(idx, imgname))
105
+
106
+ # summarize psnr/ssim
107
+ if img_gt is not None:
108
+ ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
109
+ ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
110
+ print('\n{} \n-- Average PSNR/SSIM(RGB): {:.2f} dB; {:.4f}'.format(save_dir, ave_psnr, ave_ssim))
111
+ if img_gt.ndim == 3:
112
+ ave_psnr_y = sum(test_results['psnr_y']) / len(test_results['psnr_y'])
113
+ ave_ssim_y = sum(test_results['ssim_y']) / len(test_results['ssim_y'])
114
+ print('-- Average PSNR_Y/SSIM_Y: {:.2f} dB; {:.4f}'.format(ave_psnr_y, ave_ssim_y))
115
+ if args.task in ['jpeg_car']:
116
+ ave_psnr_b = sum(test_results['psnr_b']) / len(test_results['psnr_b'])
117
+ print('-- Average PSNR_B: {:.2f} dB'.format(ave_psnr_b))
118
+
119
+
120
+ def define_model(args):
121
+ # 001 classical image sr
122
+ if args.task == 'classical_sr':
123
+ model = net(upscale=args.scale, in_chans=3, img_size=args.training_patch_size, window_size=8,
124
+ img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
125
+ mlp_ratio=2, upsampler='pixelshuffle', resi_connection='1conv')
126
+ param_key_g = 'params'
127
+
128
+ # 002 lightweight image sr
129
+ # use 'pixelshuffledirect' to save parameters
130
+ elif args.task == 'lightweight_sr':
131
+ model = net(upscale=args.scale, in_chans=3, img_size=64, window_size=8,
132
+ img_range=1., depths=[6, 6, 6, 6], embed_dim=60, num_heads=[6, 6, 6, 6],
133
+ mlp_ratio=2, upsampler='pixelshuffledirect', resi_connection='1conv')
134
+ param_key_g = 'params'
135
+
136
+ # 003 real-world image sr
137
+ elif args.task == 'real_sr':
138
+ if not args.large_model:
139
+ # use 'nearest+conv' to avoid block artifacts
140
+ model = net(upscale=4, in_chans=3, img_size=64, window_size=8,
141
+ img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
142
+ mlp_ratio=2, upsampler='nearest+conv', resi_connection='1conv')
143
+ else:
144
+ # larger model size; use '3conv' to save parameters and memory; use ema for GAN training
145
+ model = net(upscale=4, in_chans=3, img_size=64, window_size=8,
146
+ img_range=1., depths=[6, 6, 6, 6, 6, 6, 6, 6, 6], embed_dim=248,
147
+ num_heads=[8, 8, 8, 8, 8, 8, 8, 8, 8],
148
+ mlp_ratio=2, upsampler='nearest+conv', resi_connection='3conv')
149
+ param_key_g = 'params_ema'
150
+
151
+ # 004 grayscale image denoising
152
+ elif args.task == 'gray_dn':
153
+ model = net(upscale=1, in_chans=1, img_size=128, window_size=8,
154
+ img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
155
+ mlp_ratio=2, upsampler='', resi_connection='1conv')
156
+ param_key_g = 'params'
157
+
158
+ # 005 color image denoising
159
+ elif args.task == 'color_dn':
160
+ model = net(upscale=1, in_chans=3, img_size=128, window_size=8,
161
+ img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
162
+ mlp_ratio=2, upsampler='', resi_connection='1conv')
163
+ param_key_g = 'params'
164
+
165
+ # 006 JPEG compression artifact reduction
166
+ # use window_size=7 because JPEG encoding uses 8x8; use img_range=255 because it's sligtly better than 1
167
+ elif args.task == 'jpeg_car':
168
+ model = net(upscale=1, in_chans=1, img_size=126, window_size=7,
169
+ img_range=255., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
170
+ mlp_ratio=2, upsampler='', resi_connection='1conv')
171
+ param_key_g = 'params'
172
+
173
+ pretrained_model = torch.load(args.model_path)
174
+ model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True)
175
+
176
+ return model
177
+
178
+
179
+ def setup(args):
180
+ # 001 classical image sr/ 002 lightweight image sr
181
+ if args.task in ['classical_sr', 'lightweight_sr']:
182
+ save_dir = f'results/swinir_{args.task}_x{args.scale}'
183
+ folder = args.folder_gt
184
+ border = args.scale
185
+ window_size = 8
186
+
187
+ # 003 real-world image sr
188
+ elif args.task in ['real_sr']:
189
+ save_dir = f'results/swinir_{args.task}_x{args.scale}'
190
+ folder = args.folder_lq
191
+ border = 0
192
+ window_size = 8
193
+
194
+ # 004 grayscale image denoising/ 005 color image denoising
195
+ elif args.task in ['gray_dn', 'color_dn']:
196
+ save_dir = f'results/swinir_{args.task}_noise{args.noise}'
197
+ folder = args.folder_gt
198
+ border = 0
199
+ window_size = 8
200
+
201
+ # 006 JPEG compression artifact reduction
202
+ elif args.task in ['jpeg_car']:
203
+ save_dir = f'results/swinir_{args.task}_jpeg{args.jpeg}'
204
+ folder = args.folder_gt
205
+ border = 0
206
+ window_size = 7
207
+
208
+ return folder, save_dir, border, window_size
209
+
210
+
211
+ def get_image_pair(args, path):
212
+ (imgname, imgext) = os.path.splitext(os.path.basename(path))
213
+
214
+ # 001 classical image sr/ 002 lightweight image sr (load lq-gt image pairs)
215
+ if args.task in ['classical_sr', 'lightweight_sr']:
216
+ img_gt = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
217
+ img_lq = cv2.imread(f'{args.folder_lq}/{imgname}x{args.scale}{imgext}', cv2.IMREAD_COLOR).astype(
218
+ np.float32) / 255.
219
+
220
+ # 003 real-world image sr (load lq image only)
221
+ elif args.task in ['real_sr']:
222
+ img_gt = None
223
+ img_lq = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
224
+
225
+ # 004 grayscale image denoising (load gt image and generate lq image on-the-fly)
226
+ elif args.task in ['gray_dn']:
227
+ img_gt = cv2.imread(path, cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255.
228
+ np.random.seed(seed=0)
229
+ img_lq = img_gt + np.random.normal(0, args.noise / 255., img_gt.shape)
230
+ img_gt = np.expand_dims(img_gt, axis=2)
231
+ img_lq = np.expand_dims(img_lq, axis=2)
232
+
233
+ # 005 color image denoising (load gt image and generate lq image on-the-fly)
234
+ elif args.task in ['color_dn']:
235
+ img_gt = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
236
+ np.random.seed(seed=0)
237
+ img_lq = img_gt + np.random.normal(0, args.noise / 255., img_gt.shape)
238
+
239
+ # 006 JPEG compression artifact reduction (load gt image and generate lq image on-the-fly)
240
+ elif args.task in ['jpeg_car']:
241
+ img_gt = cv2.imread(path, cv2.IMREAD_UNCHANGED)
242
+ if img_gt.ndim != 2:
243
+ img_gt = util.rgb2ycbcr(cv2.cvtColor(img_gt, cv2.COLOR_BGR2RGB), y_only=True)
244
+ result, encimg = cv2.imencode('.jpg', img_gt, [int(cv2.IMWRITE_JPEG_QUALITY), args.jpeg])
245
+ img_lq = cv2.imdecode(encimg, 0)
246
+ img_gt = np.expand_dims(img_gt, axis=2).astype(np.float32) / 255.
247
+ img_lq = np.expand_dims(img_lq, axis=2).astype(np.float32) / 255.
248
+
249
+ return imgname, img_lq, img_gt
250
+
251
+
252
+ if __name__ == '__main__':
253
+ main()
network_swinir.py ADDED
@@ -0,0 +1,854 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -----------------------------------------------------------------------------------
2
+ # SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257
3
+ # Originally Written by Ze Liu, Modified by Jingyun Liang.
4
+ # -----------------------------------------------------------------------------------
5
+
6
+ import math
7
+ import torch
8
+ import torch.nn as nn
9
+ import torch.utils.checkpoint as checkpoint
10
+ from timm.models.layers import DropPath, to_2tuple, trunc_normal_
11
+
12
+
13
+ class Mlp(nn.Module):
14
+ def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
15
+ super().__init__()
16
+ out_features = out_features or in_features
17
+ hidden_features = hidden_features or in_features
18
+ self.fc1 = nn.Linear(in_features, hidden_features)
19
+ self.act = act_layer()
20
+ self.fc2 = nn.Linear(hidden_features, out_features)
21
+ self.drop = nn.Dropout(drop)
22
+
23
+ def forward(self, x):
24
+ x = self.fc1(x)
25
+ x = self.act(x)
26
+ x = self.drop(x)
27
+ x = self.fc2(x)
28
+ x = self.drop(x)
29
+ return x
30
+
31
+
32
+ def window_partition(x, window_size):
33
+ """
34
+ Args:
35
+ x: (B, H, W, C)
36
+ window_size (int): window size
37
+
38
+ Returns:
39
+ windows: (num_windows*B, window_size, window_size, C)
40
+ """
41
+ B, H, W, C = x.shape
42
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
43
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
44
+ return windows
45
+
46
+
47
+ def window_reverse(windows, window_size, H, W):
48
+ """
49
+ Args:
50
+ windows: (num_windows*B, window_size, window_size, C)
51
+ window_size (int): Window size
52
+ H (int): Height of image
53
+ W (int): Width of image
54
+
55
+ Returns:
56
+ x: (B, H, W, C)
57
+ """
58
+ B = int(windows.shape[0] / (H * W / window_size / window_size))
59
+ x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
60
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
61
+ return x
62
+
63
+
64
+ class WindowAttention(nn.Module):
65
+ r""" Window based multi-head self attention (W-MSA) module with relative position bias.
66
+ It supports both of shifted and non-shifted window.
67
+
68
+ Args:
69
+ dim (int): Number of input channels.
70
+ window_size (tuple[int]): The height and width of the window.
71
+ num_heads (int): Number of attention heads.
72
+ qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
73
+ qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
74
+ attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
75
+ proj_drop (float, optional): Dropout ratio of output. Default: 0.0
76
+ """
77
+
78
+ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
79
+
80
+ super().__init__()
81
+ self.dim = dim
82
+ self.window_size = window_size # Wh, Ww
83
+ self.num_heads = num_heads
84
+ head_dim = dim // num_heads
85
+ self.scale = qk_scale or head_dim ** -0.5
86
+
87
+ # define a parameter table of relative position bias
88
+ self.relative_position_bias_table = nn.Parameter(
89
+ torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH
90
+
91
+ # get pair-wise relative position index for each token inside the window
92
+ coords_h = torch.arange(self.window_size[0])
93
+ coords_w = torch.arange(self.window_size[1])
94
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
95
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
96
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
97
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
98
+ relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
99
+ relative_coords[:, :, 1] += self.window_size[1] - 1
100
+ relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
101
+ relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
102
+ self.register_buffer("relative_position_index", relative_position_index)
103
+
104
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
105
+ self.attn_drop = nn.Dropout(attn_drop)
106
+ self.proj = nn.Linear(dim, dim)
107
+
108
+ self.proj_drop = nn.Dropout(proj_drop)
109
+
110
+ trunc_normal_(self.relative_position_bias_table, std=.02)
111
+ self.softmax = nn.Softmax(dim=-1)
112
+
113
+ def forward(self, x, mask=None):
114
+ """
115
+ Args:
116
+ x: input features with shape of (num_windows*B, N, C)
117
+ mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
118
+ """
119
+ B_, N, C = x.shape
120
+ qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
121
+ q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
122
+
123
+ q = q * self.scale
124
+ attn = (q @ k.transpose(-2, -1))
125
+
126
+ relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
127
+ self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH
128
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
129
+ attn = attn + relative_position_bias.unsqueeze(0)
130
+
131
+ if mask is not None:
132
+ nW = mask.shape[0]
133
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
134
+ attn = attn.view(-1, self.num_heads, N, N)
135
+ attn = self.softmax(attn)
136
+ else:
137
+ attn = self.softmax(attn)
138
+
139
+ attn = self.attn_drop(attn)
140
+
141
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
142
+ x = self.proj(x)
143
+ x = self.proj_drop(x)
144
+ return x
145
+
146
+ def extra_repr(self) -> str:
147
+ return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
148
+
149
+ def flops(self, N):
150
+ # calculate flops for 1 window with token length of N
151
+ flops = 0
152
+ # qkv = self.qkv(x)
153
+ flops += N * self.dim * 3 * self.dim
154
+ # attn = (q @ k.transpose(-2, -1))
155
+ flops += self.num_heads * N * (self.dim // self.num_heads) * N
156
+ # x = (attn @ v)
157
+ flops += self.num_heads * N * N * (self.dim // self.num_heads)
158
+ # x = self.proj(x)
159
+ flops += N * self.dim * self.dim
160
+ return flops
161
+
162
+
163
+ class SwinTransformerBlock(nn.Module):
164
+ r""" Swin Transformer Block.
165
+
166
+ Args:
167
+ dim (int): Number of input channels.
168
+ input_resolution (tuple[int]): Input resulotion.
169
+ num_heads (int): Number of attention heads.
170
+ window_size (int): Window size.
171
+ shift_size (int): Shift size for SW-MSA.
172
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
173
+ qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
174
+ qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
175
+ drop (float, optional): Dropout rate. Default: 0.0
176
+ attn_drop (float, optional): Attention dropout rate. Default: 0.0
177
+ drop_path (float, optional): Stochastic depth rate. Default: 0.0
178
+ act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
179
+ norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
180
+ """
181
+
182
+ def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
183
+ mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
184
+ act_layer=nn.GELU, norm_layer=nn.LayerNorm):
185
+ super().__init__()
186
+ self.dim = dim
187
+ self.input_resolution = input_resolution
188
+ self.num_heads = num_heads
189
+ self.window_size = window_size
190
+ self.shift_size = shift_size
191
+ self.mlp_ratio = mlp_ratio
192
+ if min(self.input_resolution) <= self.window_size:
193
+ # if window size is larger than input resolution, we don't partition windows
194
+ self.shift_size = 0
195
+ self.window_size = min(self.input_resolution)
196
+ assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
197
+
198
+ self.norm1 = norm_layer(dim)
199
+ self.attn = WindowAttention(
200
+ dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
201
+ qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
202
+
203
+ self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
204
+ self.norm2 = norm_layer(dim)
205
+ mlp_hidden_dim = int(dim * mlp_ratio)
206
+ self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
207
+
208
+ if self.shift_size > 0:
209
+ attn_mask = self.calculate_mask(self.input_resolution)
210
+ else:
211
+ attn_mask = None
212
+
213
+ self.register_buffer("attn_mask", attn_mask)
214
+
215
+ def calculate_mask(self, x_size):
216
+ # calculate attention mask for SW-MSA
217
+ H, W = x_size
218
+ img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1
219
+ h_slices = (slice(0, -self.window_size),
220
+ slice(-self.window_size, -self.shift_size),
221
+ slice(-self.shift_size, None))
222
+ w_slices = (slice(0, -self.window_size),
223
+ slice(-self.window_size, -self.shift_size),
224
+ slice(-self.shift_size, None))
225
+ cnt = 0
226
+ for h in h_slices:
227
+ for w in w_slices:
228
+ img_mask[:, h, w, :] = cnt
229
+ cnt += 1
230
+
231
+ mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1
232
+ mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
233
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
234
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
235
+
236
+ return attn_mask
237
+
238
+ def forward(self, x, x_size):
239
+ H, W = x_size
240
+ B, L, C = x.shape
241
+ # assert L == H * W, "input feature has wrong size"
242
+
243
+ shortcut = x
244
+ x = self.norm1(x)
245
+ x = x.view(B, H, W, C)
246
+
247
+ # cyclic shift
248
+ if self.shift_size > 0:
249
+ shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
250
+ else:
251
+ shifted_x = x
252
+
253
+ # partition windows
254
+ x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C
255
+ x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C
256
+
257
+ # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
258
+ if self.input_resolution == x_size:
259
+ attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C
260
+ else:
261
+ attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))
262
+
263
+ # merge windows
264
+ attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
265
+ shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C
266
+
267
+ # reverse cyclic shift
268
+ if self.shift_size > 0:
269
+ x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
270
+ else:
271
+ x = shifted_x
272
+ x = x.view(B, H * W, C)
273
+
274
+ # FFN
275
+ x = shortcut + self.drop_path(x)
276
+ x = x + self.drop_path(self.mlp(self.norm2(x)))
277
+
278
+ return x
279
+
280
+ def extra_repr(self) -> str:
281
+ return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
282
+ f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
283
+
284
+ def flops(self):
285
+ flops = 0
286
+ H, W = self.input_resolution
287
+ # norm1
288
+ flops += self.dim * H * W
289
+ # W-MSA/SW-MSA
290
+ nW = H * W / self.window_size / self.window_size
291
+ flops += nW * self.attn.flops(self.window_size * self.window_size)
292
+ # mlp
293
+ flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
294
+ # norm2
295
+ flops += self.dim * H * W
296
+ return flops
297
+
298
+
299
+ class PatchMerging(nn.Module):
300
+ r""" Patch Merging Layer.
301
+
302
+ Args:
303
+ input_resolution (tuple[int]): Resolution of input feature.
304
+ dim (int): Number of input channels.
305
+ norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
306
+ """
307
+
308
+ def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
309
+ super().__init__()
310
+ self.input_resolution = input_resolution
311
+ self.dim = dim
312
+ self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
313
+ self.norm = norm_layer(4 * dim)
314
+
315
+ def forward(self, x):
316
+ """
317
+ x: B, H*W, C
318
+ """
319
+ H, W = self.input_resolution
320
+ B, L, C = x.shape
321
+ assert L == H * W, "input feature has wrong size"
322
+ assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
323
+
324
+ x = x.view(B, H, W, C)
325
+
326
+ x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C
327
+ x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C
328
+ x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C
329
+ x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C
330
+ x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C
331
+ x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C
332
+
333
+ x = self.norm(x)
334
+ x = self.reduction(x)
335
+
336
+ return x
337
+
338
+ def extra_repr(self) -> str:
339
+ return f"input_resolution={self.input_resolution}, dim={self.dim}"
340
+
341
+ def flops(self):
342
+ H, W = self.input_resolution
343
+ flops = H * W * self.dim
344
+ flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
345
+ return flops
346
+
347
+
348
+ class BasicLayer(nn.Module):
349
+ """ A basic Swin Transformer layer for one stage.
350
+
351
+ Args:
352
+ dim (int): Number of input channels.
353
+ input_resolution (tuple[int]): Input resolution.
354
+ depth (int): Number of blocks.
355
+ num_heads (int): Number of attention heads.
356
+ window_size (int): Local window size.
357
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
358
+ qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
359
+ qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
360
+ drop (float, optional): Dropout rate. Default: 0.0
361
+ attn_drop (float, optional): Attention dropout rate. Default: 0.0
362
+ drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
363
+ norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
364
+ downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
365
+ use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
366
+ """
367
+
368
+ def __init__(self, dim, input_resolution, depth, num_heads, window_size,
369
+ mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
370
+ drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):
371
+
372
+ super().__init__()
373
+ self.dim = dim
374
+ self.input_resolution = input_resolution
375
+ self.depth = depth
376
+ self.use_checkpoint = use_checkpoint
377
+
378
+ # build blocks
379
+ self.blocks = nn.ModuleList([
380
+ SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
381
+ num_heads=num_heads, window_size=window_size,
382
+ shift_size=0 if (i % 2 == 0) else window_size // 2,
383
+ mlp_ratio=mlp_ratio,
384
+ qkv_bias=qkv_bias, qk_scale=qk_scale,
385
+ drop=drop, attn_drop=attn_drop,
386
+ drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
387
+ norm_layer=norm_layer)
388
+ for i in range(depth)])
389
+
390
+ # patch merging layer
391
+ if downsample is not None:
392
+ self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
393
+ else:
394
+ self.downsample = None
395
+
396
+ def forward(self, x, x_size):
397
+ for blk in self.blocks:
398
+ if self.use_checkpoint:
399
+ x = checkpoint.checkpoint(blk, x, x_size)
400
+ else:
401
+ x = blk(x, x_size)
402
+ if self.downsample is not None:
403
+ x = self.downsample(x)
404
+ return x
405
+
406
+ def extra_repr(self) -> str:
407
+ return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
408
+
409
+ def flops(self):
410
+ flops = 0
411
+ for blk in self.blocks:
412
+ flops += blk.flops()
413
+ if self.downsample is not None:
414
+ flops += self.downsample.flops()
415
+ return flops
416
+
417
+
418
+ class RSTB(nn.Module):
419
+ """Residual Swin Transformer Block (RSTB).
420
+
421
+ Args:
422
+ dim (int): Number of input channels.
423
+ input_resolution (tuple[int]): Input resolution.
424
+ depth (int): Number of blocks.
425
+ num_heads (int): Number of attention heads.
426
+ window_size (int): Local window size.
427
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
428
+ qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
429
+ qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
430
+ drop (float, optional): Dropout rate. Default: 0.0
431
+ attn_drop (float, optional): Attention dropout rate. Default: 0.0
432
+ drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
433
+ norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
434
+ downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
435
+ use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
436
+ img_size: Input image size.
437
+ patch_size: Patch size.
438
+ resi_connection: The convolutional block before residual connection.
439
+ """
440
+
441
+ def __init__(self, dim, input_resolution, depth, num_heads, window_size,
442
+ mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
443
+ drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
444
+ img_size=224, patch_size=4, resi_connection='1conv'):
445
+ super(RSTB, self).__init__()
446
+
447
+ self.dim = dim
448
+ self.input_resolution = input_resolution
449
+
450
+ self.residual_group = BasicLayer(dim=dim,
451
+ input_resolution=input_resolution,
452
+ depth=depth,
453
+ num_heads=num_heads,
454
+ window_size=window_size,
455
+ mlp_ratio=mlp_ratio,
456
+ qkv_bias=qkv_bias, qk_scale=qk_scale,
457
+ drop=drop, attn_drop=attn_drop,
458
+ drop_path=drop_path,
459
+ norm_layer=norm_layer,
460
+ downsample=downsample,
461
+ use_checkpoint=use_checkpoint)
462
+
463
+ if resi_connection == '1conv':
464
+ self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
465
+ elif resi_connection == '3conv':
466
+ # to save parameters and memory
467
+ self.conv = nn.Sequential(nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),
468
+ nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
469
+ nn.LeakyReLU(negative_slope=0.2, inplace=True),
470
+ nn.Conv2d(dim // 4, dim, 3, 1, 1))
471
+
472
+ self.patch_embed = PatchEmbed(
473
+ img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,
474
+ norm_layer=None)
475
+
476
+ self.patch_unembed = PatchUnEmbed(
477
+ img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,
478
+ norm_layer=None)
479
+
480
+ def forward(self, x, x_size):
481
+ return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x
482
+
483
+ def flops(self):
484
+ flops = 0
485
+ flops += self.residual_group.flops()
486
+ H, W = self.input_resolution
487
+ flops += H * W * self.dim * self.dim * 9
488
+ flops += self.patch_embed.flops()
489
+ flops += self.patch_unembed.flops()
490
+
491
+ return flops
492
+
493
+
494
+ class PatchEmbed(nn.Module):
495
+ r""" Image to Patch Embedding
496
+
497
+ Args:
498
+ img_size (int): Image size. Default: 224.
499
+ patch_size (int): Patch token size. Default: 4.
500
+ in_chans (int): Number of input image channels. Default: 3.
501
+ embed_dim (int): Number of linear projection output channels. Default: 96.
502
+ norm_layer (nn.Module, optional): Normalization layer. Default: None
503
+ """
504
+
505
+ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
506
+ super().__init__()
507
+ img_size = to_2tuple(img_size)
508
+ patch_size = to_2tuple(patch_size)
509
+ patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
510
+ self.img_size = img_size
511
+ self.patch_size = patch_size
512
+ self.patches_resolution = patches_resolution
513
+ self.num_patches = patches_resolution[0] * patches_resolution[1]
514
+
515
+ self.in_chans = in_chans
516
+ self.embed_dim = embed_dim
517
+
518
+ if norm_layer is not None:
519
+ self.norm = norm_layer(embed_dim)
520
+ else:
521
+ self.norm = None
522
+
523
+ def forward(self, x):
524
+ x = x.flatten(2).transpose(1, 2) # B Ph*Pw C
525
+ if self.norm is not None:
526
+ x = self.norm(x)
527
+ return x
528
+
529
+ def flops(self):
530
+ flops = 0
531
+ H, W = self.img_size
532
+ if self.norm is not None:
533
+ flops += H * W * self.embed_dim
534
+ return flops
535
+
536
+
537
+ class PatchUnEmbed(nn.Module):
538
+ r""" Image to Patch Unembedding
539
+
540
+ Args:
541
+ img_size (int): Image size. Default: 224.
542
+ patch_size (int): Patch token size. Default: 4.
543
+ in_chans (int): Number of input image channels. Default: 3.
544
+ embed_dim (int): Number of linear projection output channels. Default: 96.
545
+ norm_layer (nn.Module, optional): Normalization layer. Default: None
546
+ """
547
+
548
+ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
549
+ super().__init__()
550
+ img_size = to_2tuple(img_size)
551
+ patch_size = to_2tuple(patch_size)
552
+ patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
553
+ self.img_size = img_size
554
+ self.patch_size = patch_size
555
+ self.patches_resolution = patches_resolution
556
+ self.num_patches = patches_resolution[0] * patches_resolution[1]
557
+
558
+ self.in_chans = in_chans
559
+ self.embed_dim = embed_dim
560
+
561
+ def forward(self, x, x_size):
562
+ B, HW, C = x.shape
563
+ x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C
564
+ return x
565
+
566
+ def flops(self):
567
+ flops = 0
568
+ return flops
569
+
570
+
571
+ class Upsample(nn.Sequential):
572
+ """Upsample module.
573
+
574
+ Args:
575
+ scale (int): Scale factor. Supported scales: 2^n and 3.
576
+ num_feat (int): Channel number of intermediate features.
577
+ """
578
+
579
+ def __init__(self, scale, num_feat):
580
+ m = []
581
+ if (scale & (scale - 1)) == 0: # scale = 2^n
582
+ for _ in range(int(math.log(scale, 2))):
583
+ m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
584
+ m.append(nn.PixelShuffle(2))
585
+ elif scale == 3:
586
+ m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
587
+ m.append(nn.PixelShuffle(3))
588
+ else:
589
+ raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
590
+ super(Upsample, self).__init__(*m)
591
+
592
+
593
+ class UpsampleOneStep(nn.Sequential):
594
+ """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
595
+ Used in lightweight SR to save parameters.
596
+
597
+ Args:
598
+ scale (int): Scale factor. Supported scales: 2^n and 3.
599
+ num_feat (int): Channel number of intermediate features.
600
+
601
+ """
602
+
603
+ def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):
604
+ self.num_feat = num_feat
605
+ self.input_resolution = input_resolution
606
+ m = []
607
+ m.append(nn.Conv2d(num_feat, (scale ** 2) * num_out_ch, 3, 1, 1))
608
+ m.append(nn.PixelShuffle(scale))
609
+ super(UpsampleOneStep, self).__init__(*m)
610
+
611
+ def flops(self):
612
+ H, W = self.input_resolution
613
+ flops = H * W * self.num_feat * 3 * 9
614
+ return flops
615
+
616
+
617
+ class SwinIR(nn.Module):
618
+ r""" SwinIR
619
+ A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer.
620
+
621
+ Args:
622
+ img_size (int | tuple(int)): Input image size. Default 64
623
+ patch_size (int | tuple(int)): Patch size. Default: 1
624
+ in_chans (int): Number of input image channels. Default: 3
625
+ embed_dim (int): Patch embedding dimension. Default: 96
626
+ depths (tuple(int)): Depth of each Swin Transformer layer.
627
+ num_heads (tuple(int)): Number of attention heads in different layers.
628
+ window_size (int): Window size. Default: 7
629
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
630
+ qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
631
+ qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
632
+ drop_rate (float): Dropout rate. Default: 0
633
+ attn_drop_rate (float): Attention dropout rate. Default: 0
634
+ drop_path_rate (float): Stochastic depth rate. Default: 0.1
635
+ norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
636
+ ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
637
+ patch_norm (bool): If True, add normalization after patch embedding. Default: True
638
+ use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
639
+ upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
640
+ img_range: Image range. 1. or 255.
641
+ upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
642
+ resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
643
+ """
644
+
645
+ def __init__(self, img_size=64, patch_size=1, in_chans=3,
646
+ embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],
647
+ window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
648
+ drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
649
+ norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
650
+ use_checkpoint=False, upscale=2, img_range=1., upsampler='', resi_connection='1conv',
651
+ **kwargs):
652
+ super(SwinIR, self).__init__()
653
+ num_in_ch = in_chans
654
+ num_out_ch = in_chans
655
+ num_feat = 64
656
+ self.img_range = img_range
657
+ if in_chans == 3:
658
+ rgb_mean = (0.4488, 0.4371, 0.4040)
659
+ self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
660
+ else:
661
+ self.mean = torch.zeros(1, 1, 1, 1)
662
+ self.upscale = upscale
663
+ self.upsampler = upsampler
664
+
665
+ #####################################################################################################
666
+ ################################### 1, shallow feature extraction ###################################
667
+ self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
668
+
669
+ #####################################################################################################
670
+ ################################### 2, deep feature extraction ######################################
671
+ self.num_layers = len(depths)
672
+ self.embed_dim = embed_dim
673
+ self.ape = ape
674
+ self.patch_norm = patch_norm
675
+ self.num_features = embed_dim
676
+ self.mlp_ratio = mlp_ratio
677
+
678
+ # split image into non-overlapping patches
679
+ self.patch_embed = PatchEmbed(
680
+ img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
681
+ norm_layer=norm_layer if self.patch_norm else None)
682
+ num_patches = self.patch_embed.num_patches
683
+ patches_resolution = self.patch_embed.patches_resolution
684
+ self.patches_resolution = patches_resolution
685
+
686
+ # merge non-overlapping patches into image
687
+ self.patch_unembed = PatchUnEmbed(
688
+ img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
689
+ norm_layer=norm_layer if self.patch_norm else None)
690
+
691
+ # absolute position embedding
692
+ if self.ape:
693
+ self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
694
+ trunc_normal_(self.absolute_pos_embed, std=.02)
695
+
696
+ self.pos_drop = nn.Dropout(p=drop_rate)
697
+
698
+ # stochastic depth
699
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule
700
+
701
+ # build Residual Swin Transformer blocks (RSTB)
702
+ self.layers = nn.ModuleList()
703
+ for i_layer in range(self.num_layers):
704
+ layer = RSTB(dim=embed_dim,
705
+ input_resolution=(patches_resolution[0],
706
+ patches_resolution[1]),
707
+ depth=depths[i_layer],
708
+ num_heads=num_heads[i_layer],
709
+ window_size=window_size,
710
+ mlp_ratio=self.mlp_ratio,
711
+ qkv_bias=qkv_bias, qk_scale=qk_scale,
712
+ drop=drop_rate, attn_drop=attn_drop_rate,
713
+ drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], # no impact on SR results
714
+ norm_layer=norm_layer,
715
+ downsample=None,
716
+ use_checkpoint=use_checkpoint,
717
+ img_size=img_size,
718
+ patch_size=patch_size,
719
+ resi_connection=resi_connection
720
+
721
+ )
722
+ self.layers.append(layer)
723
+ self.norm = norm_layer(self.num_features)
724
+
725
+ # build the last conv layer in deep feature extraction
726
+ if resi_connection == '1conv':
727
+ self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
728
+ elif resi_connection == '3conv':
729
+ # to save parameters and memory
730
+ self.conv_after_body = nn.Sequential(nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
731
+ nn.LeakyReLU(negative_slope=0.2, inplace=True),
732
+ nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
733
+ nn.LeakyReLU(negative_slope=0.2, inplace=True),
734
+ nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1))
735
+
736
+ #####################################################################################################
737
+ ################################ 3, high quality image reconstruction ################################
738
+ if self.upsampler == 'pixelshuffle':
739
+ # for classical SR
740
+ self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
741
+ nn.LeakyReLU(inplace=True))
742
+ self.upsample = Upsample(upscale, num_feat)
743
+ self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
744
+ elif self.upsampler == 'pixelshuffledirect':
745
+ # for lightweight SR (to save parameters)
746
+ self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch,
747
+ (patches_resolution[0], patches_resolution[1]))
748
+ elif self.upsampler == 'nearest+conv':
749
+ # for real-world SR (less artifacts)
750
+ assert self.upscale == 4, 'only support x4 now.'
751
+ self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
752
+ nn.LeakyReLU(inplace=True))
753
+ self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
754
+ self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
755
+ self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
756
+ self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
757
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
758
+ else:
759
+ # for image denoising and JPEG compression artifact reduction
760
+ self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)
761
+
762
+ self.apply(self._init_weights)
763
+
764
+ def _init_weights(self, m):
765
+ if isinstance(m, nn.Linear):
766
+ trunc_normal_(m.weight, std=.02)
767
+ if isinstance(m, nn.Linear) and m.bias is not None:
768
+ nn.init.constant_(m.bias, 0)
769
+ elif isinstance(m, nn.LayerNorm):
770
+ nn.init.constant_(m.bias, 0)
771
+ nn.init.constant_(m.weight, 1.0)
772
+
773
+ @torch.jit.ignore
774
+ def no_weight_decay(self):
775
+ return {'absolute_pos_embed'}
776
+
777
+ @torch.jit.ignore
778
+ def no_weight_decay_keywords(self):
779
+ return {'relative_position_bias_table'}
780
+
781
+ def forward_features(self, x):
782
+ x_size = (x.shape[2], x.shape[3])
783
+ x = self.patch_embed(x)
784
+ if self.ape:
785
+ x = x + self.absolute_pos_embed
786
+ x = self.pos_drop(x)
787
+
788
+ for layer in self.layers:
789
+ x = layer(x, x_size)
790
+
791
+ x = self.norm(x) # B L C
792
+ x = self.patch_unembed(x, x_size)
793
+
794
+ return x
795
+
796
+ def forward(self, x):
797
+ self.mean = self.mean.type_as(x)
798
+ x = (x - self.mean) * self.img_range
799
+
800
+ if self.upsampler == 'pixelshuffle':
801
+ # for classical SR
802
+ x = self.conv_first(x)
803
+ x = self.conv_after_body(self.forward_features(x)) + x
804
+ x = self.conv_before_upsample(x)
805
+ x = self.conv_last(self.upsample(x))
806
+ elif self.upsampler == 'pixelshuffledirect':
807
+ # for lightweight SR
808
+ x = self.conv_first(x)
809
+ x = self.conv_after_body(self.forward_features(x)) + x
810
+ x = self.upsample(x)
811
+ elif self.upsampler == 'nearest+conv':
812
+ # for real-world SR
813
+ x = self.conv_first(x)
814
+ x = self.conv_after_body(self.forward_features(x)) + x
815
+ x = self.conv_before_upsample(x)
816
+ x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
817
+ x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
818
+ x = self.conv_last(self.lrelu(self.conv_hr(x)))
819
+ else:
820
+ # for image denoising and JPEG compression artifact reduction
821
+ x_first = self.conv_first(x)
822
+ res = self.conv_after_body(self.forward_features(x_first)) + x_first
823
+ x = x + self.conv_last(res)
824
+
825
+ x = x / self.img_range + self.mean
826
+
827
+ return x
828
+
829
+ def flops(self):
830
+ flops = 0
831
+ H, W = self.patches_resolution
832
+ flops += H * W * 3 * self.embed_dim * 9
833
+ flops += self.patch_embed.flops()
834
+ for i, layer in enumerate(self.layers):
835
+ flops += layer.flops()
836
+ flops += H * W * 3 * self.embed_dim * self.embed_dim
837
+ flops += self.upsample.flops()
838
+ return flops
839
+
840
+
841
+ if __name__ == '__main__':
842
+ upscale = 4
843
+ window_size = 8
844
+ height = (1024 // upscale // window_size + 1) * window_size
845
+ width = (720 // upscale // window_size + 1) * window_size
846
+ model = SwinIR(upscale=2, img_size=(height, width),
847
+ window_size=window_size, img_range=1., depths=[6, 6, 6, 6],
848
+ embed_dim=60, num_heads=[6, 6, 6, 6], mlp_ratio=2, upsampler='pixelshuffledirect')
849
+ print(model)
850
+ print(height, width, model.flops() / 1e9)
851
+
852
+ x = torch.randn((1, 3, height, width))
853
+ x = model(x)
854
+ print(x.shape)
predict.py ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cog
2
+ import tempfile
3
+ from pathlib import Path
4
+ import argparse
5
+ import shutil
6
+ import os
7
+ import cv2
8
+ import glob
9
+ import torch
10
+ from collections import OrderedDict
11
+ import numpy as np
12
+ from main_test_swinir import define_model, setup, get_image_pair
13
+
14
+
15
+ class Predictor(cog.Predictor):
16
+ def setup(self):
17
+ model_dir = 'experiments/pretrained_models'
18
+
19
+ self.model_zoo = {
20
+ 'real_sr': {
21
+ 4: os.path.join(model_dir, '003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth')
22
+ },
23
+ 'gray_dn': {
24
+ 15: os.path.join(model_dir, '004_grayDN_DFWB_s128w8_SwinIR-M_noise15.pth'),
25
+ 25: os.path.join(model_dir, '004_grayDN_DFWB_s128w8_SwinIR-M_noise25.pth'),
26
+ 50: os.path.join(model_dir, '004_grayDN_DFWB_s128w8_SwinIR-M_noise50.pth')
27
+ },
28
+ 'color_dn': {
29
+ 15: os.path.join(model_dir, '005_colorDN_DFWB_s128w8_SwinIR-M_noise15.pth'),
30
+ 25: os.path.join(model_dir, '005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth'),
31
+ 50: os.path.join(model_dir, '005_colorDN_DFWB_s128w8_SwinIR-M_noise50.pth')
32
+ },
33
+ 'jpeg_car': {
34
+ 10: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg10.pth'),
35
+ 20: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg20.pth'),
36
+ 30: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg30.pth'),
37
+ 40: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg40.pth')
38
+ }
39
+ }
40
+
41
+ parser = argparse.ArgumentParser()
42
+ parser.add_argument('--task', type=str, default='real_sr', help='classical_sr, lightweight_sr, real_sr, '
43
+ 'gray_dn, color_dn, jpeg_car')
44
+ parser.add_argument('--scale', type=int, default=1, help='scale factor: 1, 2, 3, 4, 8') # 1 for dn and jpeg car
45
+ parser.add_argument('--noise', type=int, default=15, help='noise level: 15, 25, 50')
46
+ parser.add_argument('--jpeg', type=int, default=40, help='scale factor: 10, 20, 30, 40')
47
+ parser.add_argument('--training_patch_size', type=int, default=128, help='patch size used in training SwinIR. '
48
+ 'Just used to differentiate two different settings in Table 2 of the paper. '
49
+ 'Images are NOT tested patch by patch.')
50
+ parser.add_argument('--large_model', action='store_true',
51
+ help='use large model, only provided for real image sr')
52
+ parser.add_argument('--model_path', type=str,
53
+ default=self.model_zoo['real_sr'][4])
54
+ parser.add_argument('--folder_lq', type=str, default=None, help='input low-quality test image folder')
55
+ parser.add_argument('--folder_gt', type=str, default=None, help='input ground-truth test image folder')
56
+
57
+ self.args = parser.parse_args('')
58
+
59
+ self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
60
+
61
+ self.tasks = {
62
+ 'Real-World Image Super-Resolution': 'real_sr',
63
+ 'Grayscale Image Denoising': 'gray_dn',
64
+ 'Color Image Denoising': 'color_dn',
65
+ 'JPEG Compression Artifact Reduction': 'jpeg_car'
66
+ }
67
+
68
+ @cog.input("image", type=Path, help="input image")
69
+ @cog.input("task_type", type=str, default='Real-World Image Super-Resolution',
70
+ options=['Real-World Image Super-Resolution', 'Grayscale Image Denoising', 'Color Image Denoising',
71
+ 'JPEG Compression Artifact Reduction'],
72
+ help="image restoration task type")
73
+ @cog.input("noise", type=int, default=15, options=[15, 25, 50],
74
+ help='noise level, activated for Grayscale Image Denoising and Color Image Denoising. '
75
+ 'Leave it as default or arbitrary if other tasks are selected')
76
+ @cog.input("jpeg", type=int, default=40, options=[10, 20, 30, 40],
77
+ help='scale factor, activated for JPEG Compression Artifact Reduction. '
78
+ 'Leave it as default or arbitrary if other tasks are selected')
79
+ def predict(self, image, task_type='Real-World Image Super-Resolution', jpeg=40, noise=15):
80
+
81
+ self.args.task = self.tasks[task_type]
82
+ self.args.noise = noise
83
+ self.args.jpeg = jpeg
84
+
85
+ # set model path
86
+ if self.args.task == 'real_sr':
87
+ self.args.scale = 4
88
+ self.args.model_path = self.model_zoo[self.args.task][4]
89
+ elif self.args.task in ['gray_dn', 'color_dn']:
90
+ self.args.model_path = self.model_zoo[self.args.task][noise]
91
+ else:
92
+ self.args.model_path = self.model_zoo[self.args.task][jpeg]
93
+
94
+ # set input folder
95
+ input_dir = 'input_cog_temp'
96
+ os.makedirs(input_dir, exist_ok=True)
97
+ input_path = os.path.join(input_dir, os.path.basename(image))
98
+ shutil.copy(str(image), input_path)
99
+ if self.args.task == 'real_sr':
100
+ self.args.folder_lq = input_dir
101
+ else:
102
+ self.args.folder_gt = input_dir
103
+
104
+ model = define_model(self.args)
105
+ model.eval()
106
+ model = model.to(self.device)
107
+
108
+ # setup folder and path
109
+ folder, save_dir, border, window_size = setup(self.args)
110
+ os.makedirs(save_dir, exist_ok=True)
111
+ test_results = OrderedDict()
112
+ test_results['psnr'] = []
113
+ test_results['ssim'] = []
114
+ test_results['psnr_y'] = []
115
+ test_results['ssim_y'] = []
116
+ test_results['psnr_b'] = []
117
+ # psnr, ssim, psnr_y, ssim_y, psnr_b = 0, 0, 0, 0, 0
118
+ out_path = Path(tempfile.mkdtemp()) / "out.png"
119
+
120
+ for idx, path in enumerate(sorted(glob.glob(os.path.join(folder, '*')))):
121
+ # read image
122
+ imgname, img_lq, img_gt = get_image_pair(self.args, path) # image to HWC-BGR, float32
123
+ img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]],
124
+ (2, 0, 1)) # HCW-BGR to CHW-RGB
125
+ img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(self.device) # CHW-RGB to NCHW-RGB
126
+
127
+ # inference
128
+ with torch.no_grad():
129
+ # pad input image to be a multiple of window_size
130
+ _, _, h_old, w_old = img_lq.size()
131
+ h_pad = (h_old // window_size + 1) * window_size - h_old
132
+ w_pad = (w_old // window_size + 1) * window_size - w_old
133
+ img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
134
+ img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
135
+ output = model(img_lq)
136
+ output = output[..., :h_old * self.args.scale, :w_old * self.args.scale]
137
+
138
+ # save image
139
+ output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
140
+ if output.ndim == 3:
141
+ output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR
142
+ output = (output * 255.0).round().astype(np.uint8) # float32 to uint8
143
+ cv2.imwrite(str(out_path), output)
144
+
145
+ clean_folder(input_dir)
146
+ return out_path
147
+
148
+
149
+ def clean_folder(folder):
150
+ for filename in os.listdir(folder):
151
+ file_path = os.path.join(folder, filename)
152
+ try:
153
+ if os.path.isfile(file_path) or os.path.islink(file_path):
154
+ os.unlink(file_path)
155
+ elif os.path.isdir(file_path):
156
+ shutil.rmtree(file_path)
157
+ except Exception as e:
158
+ print('Failed to delete %s. Reason: %s' % (file_path, e))
util_calculate_psnr_ssim.py ADDED
@@ -0,0 +1,346 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import torch
4
+
5
+
6
+ def calculate_psnr(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
7
+ """Calculate PSNR (Peak Signal-to-Noise Ratio).
8
+
9
+ Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
10
+
11
+ Args:
12
+ img1 (ndarray): Images with range [0, 255].
13
+ img2 (ndarray): Images with range [0, 255].
14
+ crop_border (int): Cropped pixels in each edge of an image. These
15
+ pixels are not involved in the PSNR calculation.
16
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
17
+ Default: 'HWC'.
18
+ test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
19
+
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
+
49
+ It is called by func:`calculate_ssim`.
50
+
51
+ Args:
52
+ img1 (ndarray): Images with range [0, 255] with order 'HWC'.
53
+ img2 (ndarray): Images with range [0, 255] with order 'HWC'.
54
+
55
+ Returns:
56
+ float: ssim result.
57
+ """
58
+
59
+ C1 = (0.01 * 255) ** 2
60
+ C2 = (0.03 * 255) ** 2
61
+
62
+ img1 = img1.astype(np.float64)
63
+ img2 = img2.astype(np.float64)
64
+ kernel = cv2.getGaussianKernel(11, 1.5)
65
+ window = np.outer(kernel, kernel.transpose())
66
+
67
+ mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
68
+ mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
69
+ mu1_sq = mu1 ** 2
70
+ mu2_sq = mu2 ** 2
71
+ mu1_mu2 = mu1 * mu2
72
+ sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
73
+ sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
74
+ sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
75
+
76
+ ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
77
+ return ssim_map.mean()
78
+
79
+
80
+ def calculate_ssim(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
81
+ """Calculate SSIM (structural similarity).
82
+
83
+ Ref:
84
+ Image quality assessment: From error visibility to structural similarity
85
+
86
+ The results are the same as that of the official released MATLAB code in
87
+ https://ece.uwaterloo.ca/~z70wang/research/ssim/.
88
+
89
+ For three-channel images, SSIM is calculated for each channel and then
90
+ averaged.
91
+
92
+ Args:
93
+ img1 (ndarray): Images with range [0, 255].
94
+ img2 (ndarray): Images with range [0, 255].
95
+ crop_border (int): Cropped pixels in each edge of an image. These
96
+ pixels are not involved in the SSIM calculation.
97
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
98
+ Default: 'HWC'.
99
+ test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
100
+
101
+ Returns:
102
+ float: ssim result.
103
+ """
104
+
105
+ assert img1.shape == img2.shape, (f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
106
+ if input_order not in ['HWC', 'CHW']:
107
+ raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
108
+ img1 = reorder_image(img1, input_order=input_order)
109
+ img2 = reorder_image(img2, input_order=input_order)
110
+ img1 = img1.astype(np.float64)
111
+ img2 = img2.astype(np.float64)
112
+
113
+ if crop_border != 0:
114
+ img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
115
+ img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
116
+
117
+ if test_y_channel:
118
+ img1 = to_y_channel(img1)
119
+ img2 = to_y_channel(img2)
120
+
121
+ ssims = []
122
+ for i in range(img1.shape[2]):
123
+ ssims.append(_ssim(img1[..., i], img2[..., i]))
124
+ return np.array(ssims).mean()
125
+
126
+
127
+ def _blocking_effect_factor(im):
128
+ block_size = 8
129
+
130
+ block_horizontal_positions = torch.arange(7, im.shape[3] - 1, 8)
131
+ block_vertical_positions = torch.arange(7, im.shape[2] - 1, 8)
132
+
133
+ horizontal_block_difference = (
134
+ (im[:, :, :, block_horizontal_positions] - im[:, :, :, block_horizontal_positions + 1]) ** 2).sum(
135
+ 3).sum(2).sum(1)
136
+ vertical_block_difference = (
137
+ (im[:, :, block_vertical_positions, :] - im[:, :, block_vertical_positions + 1, :]) ** 2).sum(3).sum(
138
+ 2).sum(1)
139
+
140
+ nonblock_horizontal_positions = np.setdiff1d(torch.arange(0, im.shape[3] - 1), block_horizontal_positions)
141
+ nonblock_vertical_positions = np.setdiff1d(torch.arange(0, im.shape[2] - 1), block_vertical_positions)
142
+
143
+ horizontal_nonblock_difference = (
144
+ (im[:, :, :, nonblock_horizontal_positions] - im[:, :, :, nonblock_horizontal_positions + 1]) ** 2).sum(
145
+ 3).sum(2).sum(1)
146
+ vertical_nonblock_difference = (
147
+ (im[:, :, nonblock_vertical_positions, :] - im[:, :, nonblock_vertical_positions + 1, :]) ** 2).sum(
148
+ 3).sum(2).sum(1)
149
+
150
+ n_boundary_horiz = im.shape[2] * (im.shape[3] // block_size - 1)
151
+ n_boundary_vert = im.shape[3] * (im.shape[2] // block_size - 1)
152
+ boundary_difference = (horizontal_block_difference + vertical_block_difference) / (
153
+ n_boundary_horiz + n_boundary_vert)
154
+
155
+ n_nonboundary_horiz = im.shape[2] * (im.shape[3] - 1) - n_boundary_horiz
156
+ n_nonboundary_vert = im.shape[3] * (im.shape[2] - 1) - n_boundary_vert
157
+ nonboundary_difference = (horizontal_nonblock_difference + vertical_nonblock_difference) / (
158
+ n_nonboundary_horiz + n_nonboundary_vert)
159
+
160
+ scaler = np.log2(block_size) / np.log2(min([im.shape[2], im.shape[3]]))
161
+ bef = scaler * (boundary_difference - nonboundary_difference)
162
+
163
+ bef[boundary_difference <= nonboundary_difference] = 0
164
+ return bef
165
+
166
+
167
+ def calculate_psnrb(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
168
+ """Calculate PSNR-B (Peak Signal-to-Noise Ratio).
169
+
170
+ Ref: Quality assessment of deblocked images, for JPEG image deblocking evaluation
171
+ # https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py
172
+
173
+ Args:
174
+ img1 (ndarray): Images with range [0, 255].
175
+ img2 (ndarray): Images with range [0, 255].
176
+ crop_border (int): Cropped pixels in each edge of an image. These
177
+ pixels are not involved in the PSNR calculation.
178
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
179
+ Default: 'HWC'.
180
+ test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
181
+
182
+ Returns:
183
+ float: psnr result.
184
+ """
185
+
186
+ assert img1.shape == img2.shape, (f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
187
+ if input_order not in ['HWC', 'CHW']:
188
+ raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
189
+ img1 = reorder_image(img1, input_order=input_order)
190
+ img2 = reorder_image(img2, input_order=input_order)
191
+ img1 = img1.astype(np.float64)
192
+ img2 = img2.astype(np.float64)
193
+
194
+ if crop_border != 0:
195
+ img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
196
+ img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
197
+
198
+ if test_y_channel:
199
+ img1 = to_y_channel(img1)
200
+ img2 = to_y_channel(img2)
201
+
202
+ # follow https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py
203
+ img1 = torch.from_numpy(img1).permute(2, 0, 1).unsqueeze(0) / 255.
204
+ img2 = torch.from_numpy(img2).permute(2, 0, 1).unsqueeze(0) / 255.
205
+
206
+ total = 0
207
+ for c in range(img1.shape[1]):
208
+ mse = torch.nn.functional.mse_loss(img1[:, c:c + 1, :, :], img2[:, c:c + 1, :, :], reduction='none')
209
+ bef = _blocking_effect_factor(img1[:, c:c + 1, :, :])
210
+
211
+ mse = mse.view(mse.shape[0], -1).mean(1)
212
+ total += 10 * torch.log10(1 / (mse + bef))
213
+
214
+ return float(total) / img1.shape[1]
215
+
216
+
217
+ def reorder_image(img, input_order='HWC'):
218
+ """Reorder images to 'HWC' order.
219
+
220
+ If the input_order is (h, w), return (h, w, 1);
221
+ If the input_order is (c, h, w), return (h, w, c);
222
+ If the input_order is (h, w, c), return as it is.
223
+
224
+ Args:
225
+ img (ndarray): Input image.
226
+ input_order (str): Whether the input order is 'HWC' or 'CHW'.
227
+ If the input image shape is (h, w), input_order will not have
228
+ effects. Default: 'HWC'.
229
+
230
+ Returns:
231
+ ndarray: reordered image.
232
+ """
233
+
234
+ if input_order not in ['HWC', 'CHW']:
235
+ raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' "'HWC' and 'CHW'")
236
+ if len(img.shape) == 2:
237
+ img = img[..., None]
238
+ if input_order == 'CHW':
239
+ img = img.transpose(1, 2, 0)
240
+ return img
241
+
242
+
243
+ def to_y_channel(img):
244
+ """Change to Y channel of YCbCr.
245
+
246
+ Args:
247
+ img (ndarray): Images with range [0, 255].
248
+
249
+ Returns:
250
+ (ndarray): Images with range [0, 255] (float type) without round.
251
+ """
252
+ img = img.astype(np.float32) / 255.
253
+ if img.ndim == 3 and img.shape[2] == 3:
254
+ img = bgr2ycbcr(img, y_only=True)
255
+ img = img[..., None]
256
+ return img * 255.
257
+
258
+
259
+ def _convert_input_type_range(img):
260
+ """Convert the type and range of the input image.
261
+
262
+ It converts the input image to np.float32 type and range of [0, 1].
263
+ It is mainly used for pre-processing the input image in colorspace
264
+ convertion functions such as rgb2ycbcr and ycbcr2rgb.
265
+
266
+ Args:
267
+ img (ndarray): The input image. It accepts:
268
+ 1. np.uint8 type with range [0, 255];
269
+ 2. np.float32 type with range [0, 1].
270
+
271
+ Returns:
272
+ (ndarray): The converted image with type of np.float32 and range of
273
+ [0, 1].
274
+ """
275
+ img_type = img.dtype
276
+ img = img.astype(np.float32)
277
+ if img_type == np.float32:
278
+ pass
279
+ elif img_type == np.uint8:
280
+ img /= 255.
281
+ else:
282
+ raise TypeError('The img type should be np.float32 or np.uint8, ' f'but got {img_type}')
283
+ return img
284
+
285
+
286
+ def _convert_output_type_range(img, dst_type):
287
+ """Convert the type and range of the image according to dst_type.
288
+
289
+ It converts the image to desired type and range. If `dst_type` is np.uint8,
290
+ images will be converted to np.uint8 type with range [0, 255]. If
291
+ `dst_type` is np.float32, it converts the image to np.float32 type with
292
+ range [0, 1].
293
+ It is mainly used for post-processing images in colorspace convertion
294
+ functions such as rgb2ycbcr and ycbcr2rgb.
295
+
296
+ Args:
297
+ img (ndarray): The image to be converted with np.float32 type and
298
+ range [0, 255].
299
+ dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
300
+ converts the image to np.uint8 type with range [0, 255]. If
301
+ dst_type is np.float32, it converts the image to np.float32 type
302
+ with range [0, 1].
303
+
304
+ Returns:
305
+ (ndarray): The converted image with desired type and range.
306
+ """
307
+ if dst_type not in (np.uint8, np.float32):
308
+ raise TypeError('The dst_type should be np.float32 or np.uint8, ' f'but got {dst_type}')
309
+ if dst_type == np.uint8:
310
+ img = img.round()
311
+ else:
312
+ img /= 255.
313
+ return img.astype(dst_type)
314
+
315
+
316
+ def bgr2ycbcr(img, y_only=False):
317
+ """Convert a BGR image to YCbCr image.
318
+
319
+ The bgr version of rgb2ycbcr.
320
+ It implements the ITU-R BT.601 conversion for standard-definition
321
+ television. See more details in
322
+ https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
323
+
324
+ It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
325
+ In OpenCV, it implements a JPEG conversion. See more details in
326
+ https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
327
+
328
+ Args:
329
+ img (ndarray): The input image. It accepts:
330
+ 1. np.uint8 type with range [0, 255];
331
+ 2. np.float32 type with range [0, 1].
332
+ y_only (bool): Whether to only return Y channel. Default: False.
333
+
334
+ Returns:
335
+ ndarray: The converted YCbCr image. The output image has the same type
336
+ and range as input image.
337
+ """
338
+ img_type = img.dtype
339
+ img = _convert_input_type_range(img)
340
+ if y_only:
341
+ out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
342
+ else:
343
+ out_img = np.matmul(
344
+ img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]
345
+ out_img = _convert_output_type_range(out_img, img_type)
346
+ return out_img