InvokeAI 2.2.4.5__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (169) hide show
  1. backend/__init__.py +0 -0
  2. backend/invoke_ai_web_server.py +1724 -0
  3. backend/modules/__init__.py +0 -0
  4. backend/modules/create_cmd_parser.py +55 -0
  5. backend/modules/get_canvas_generation_mode.py +117 -0
  6. backend/modules/parameters.py +71 -0
  7. backend/modules/parse_seed_weights.py +47 -0
  8. configs/INITIAL_MODELS.yaml +80 -0
  9. configs/models.yaml.example +29 -0
  10. configs/sd-concepts.txt +803 -0
  11. configs/stable-diffusion/v1-finetune.yaml +110 -0
  12. configs/stable-diffusion/v1-finetune_style.yaml +103 -0
  13. configs/stable-diffusion/v1-inference.yaml +79 -0
  14. configs/stable-diffusion/v1-inpainting-inference.yaml +79 -0
  15. configs/stable-diffusion/v1-m1-finetune.yaml +110 -0
  16. docs/assets/invoke-web-server-1.png +0 -0
  17. docs/assets/invoke_ai_banner.png +0 -0
  18. invokeai-2.2.4.5.data/data/assets/caution.png +0 -0
  19. invokeai-2.2.4.5.data/data/configs/INITIAL_MODELS.yaml +80 -0
  20. invokeai-2.2.4.5.data/data/configs/models.yaml.example +29 -0
  21. invokeai-2.2.4.5.data/data/configs/sd-concepts.txt +803 -0
  22. invokeai-2.2.4.5.data/data/configs/stable-diffusion/v1-finetune.yaml +110 -0
  23. invokeai-2.2.4.5.data/data/configs/stable-diffusion/v1-finetune_style.yaml +103 -0
  24. invokeai-2.2.4.5.data/data/configs/stable-diffusion/v1-inference.yaml +79 -0
  25. invokeai-2.2.4.5.data/data/configs/stable-diffusion/v1-inpainting-inference.yaml +79 -0
  26. invokeai-2.2.4.5.data/data/configs/stable-diffusion/v1-m1-finetune.yaml +110 -0
  27. invokeai-2.2.4.5.data/data/frontend/dist/assets/Inter-Bold.790c108b.ttf +0 -0
  28. invokeai-2.2.4.5.data/data/frontend/dist/assets/Inter.b9a8e5e2.ttf +0 -0
  29. invokeai-2.2.4.5.data/data/frontend/dist/assets/favicon.0d253ced.ico +0 -0
  30. invokeai-2.2.4.5.data/data/frontend/dist/assets/index-legacy-b741077f.js +48 -0
  31. invokeai-2.2.4.5.data/data/frontend/dist/assets/index.0dadf5d0.css +1 -0
  32. invokeai-2.2.4.5.data/data/frontend/dist/assets/index.21ff81eb.js +625 -0
  33. invokeai-2.2.4.5.data/data/frontend/dist/assets/logo.13003d72.png +0 -0
  34. invokeai-2.2.4.5.data/data/frontend/dist/assets/polyfills-legacy-dde3a68a.js +4 -0
  35. invokeai-2.2.4.5.data/data/frontend/dist/assets/polyfills.1ff60148.js +1 -0
  36. invokeai-2.2.4.5.data/data/frontend/dist/index.html +23 -0
  37. invokeai-2.2.4.5.data/data/scripts/configure_invokeai.py +797 -0
  38. invokeai-2.2.4.5.data/data/scripts/create_root.py +43 -0
  39. invokeai-2.2.4.5.data/data/scripts/dream.py +10 -0
  40. invokeai-2.2.4.5.data/data/scripts/images2prompt.py +30 -0
  41. invokeai-2.2.4.5.data/data/scripts/invoke.py +5 -0
  42. invokeai-2.2.4.5.data/data/scripts/legacy_api.py +681 -0
  43. invokeai-2.2.4.5.data/data/scripts/load_models.py +12 -0
  44. invokeai-2.2.4.5.data/data/scripts/merge_embeddings.py +124 -0
  45. invokeai-2.2.4.5.data/data/scripts/preload_models.py +12 -0
  46. invokeai-2.2.4.5.data/data/scripts/sd-metadata.py +22 -0
  47. invokeai-2.2.4.5.dist-info/LICENSE +21 -0
  48. invokeai-2.2.4.5.dist-info/LICENSE-ModelWeights.txt +294 -0
  49. invokeai-2.2.4.5.dist-info/METADATA +288 -0
  50. invokeai-2.2.4.5.dist-info/RECORD +169 -0
  51. invokeai-2.2.4.5.dist-info/WHEEL +5 -0
  52. invokeai-2.2.4.5.dist-info/entry_points.txt +4 -0
  53. invokeai-2.2.4.5.dist-info/top_level.txt +6 -0
  54. invokeai-2.2.4.5.dist-info/zip-safe +1 -0
  55. ldm/__init__.py +0 -0
  56. ldm/data/__init__.py +0 -0
  57. ldm/data/base.py +31 -0
  58. ldm/data/imagenet.py +466 -0
  59. ldm/data/lsun.py +126 -0
  60. ldm/data/personalized.py +202 -0
  61. ldm/data/personalized_style.py +169 -0
  62. ldm/generate.py +1159 -0
  63. ldm/invoke/CLI.py +960 -0
  64. ldm/invoke/__init__.py +3 -0
  65. ldm/invoke/args.py +1257 -0
  66. ldm/invoke/concepts_lib.py +169 -0
  67. ldm/invoke/conditioning.py +289 -0
  68. ldm/invoke/devices.py +27 -0
  69. ldm/invoke/generator/__init__.py +4 -0
  70. ldm/invoke/generator/base.py +335 -0
  71. ldm/invoke/generator/embiggen.py +501 -0
  72. ldm/invoke/generator/img2img.py +97 -0
  73. ldm/invoke/generator/inpaint.py +358 -0
  74. ldm/invoke/generator/omnibus.py +175 -0
  75. ldm/invoke/generator/txt2img.py +88 -0
  76. ldm/invoke/generator/txt2img2img.py +182 -0
  77. ldm/invoke/globals.py +31 -0
  78. ldm/invoke/image_util.py +70 -0
  79. ldm/invoke/log.py +66 -0
  80. ldm/invoke/model_cache.py +449 -0
  81. ldm/invoke/patchmatch.py +44 -0
  82. ldm/invoke/pngwriter.py +113 -0
  83. ldm/invoke/prompt_parser.py +655 -0
  84. ldm/invoke/readline.py +442 -0
  85. ldm/invoke/restoration/__init__.py +4 -0
  86. ldm/invoke/restoration/base.py +38 -0
  87. ldm/invoke/restoration/codeformer.py +103 -0
  88. ldm/invoke/restoration/codeformer_arch.py +276 -0
  89. ldm/invoke/restoration/gfpgan.py +84 -0
  90. ldm/invoke/restoration/outcrop.py +108 -0
  91. ldm/invoke/restoration/outpaint.py +92 -0
  92. ldm/invoke/restoration/realesrgan.py +90 -0
  93. ldm/invoke/restoration/vqgan_arch.py +435 -0
  94. ldm/invoke/seamless.py +30 -0
  95. ldm/invoke/server.py +282 -0
  96. ldm/invoke/server_legacy.py +246 -0
  97. ldm/invoke/txt2mask.py +136 -0
  98. ldm/lr_scheduler.py +143 -0
  99. ldm/models/__init__.py +0 -0
  100. ldm/models/autoencoder.py +596 -0
  101. ldm/models/diffusion/__init__.py +0 -0
  102. ldm/models/diffusion/classifier.py +355 -0
  103. ldm/models/diffusion/cross_attention_control.py +443 -0
  104. ldm/models/diffusion/cross_attention_map_saving.py +95 -0
  105. ldm/models/diffusion/ddim.py +111 -0
  106. ldm/models/diffusion/ddpm.py +2263 -0
  107. ldm/models/diffusion/ksampler.py +312 -0
  108. ldm/models/diffusion/plms.py +146 -0
  109. ldm/models/diffusion/sampler.py +450 -0
  110. ldm/models/diffusion/shared_invokeai_diffusion.py +244 -0
  111. ldm/modules/__init__.py +0 -0
  112. ldm/modules/attention.py +274 -0
  113. ldm/modules/diffusionmodules/__init__.py +0 -0
  114. ldm/modules/diffusionmodules/model.py +894 -0
  115. ldm/modules/diffusionmodules/openaimodel.py +1036 -0
  116. ldm/modules/diffusionmodules/util.py +314 -0
  117. ldm/modules/distributions/__init__.py +0 -0
  118. ldm/modules/distributions/distributions.py +104 -0
  119. ldm/modules/ema.py +88 -0
  120. ldm/modules/embedding_manager.py +413 -0
  121. ldm/modules/encoders/__init__.py +0 -0
  122. ldm/modules/encoders/modules.py +772 -0
  123. ldm/modules/image_degradation/__init__.py +6 -0
  124. ldm/modules/image_degradation/bsrgan.py +846 -0
  125. ldm/modules/image_degradation/bsrgan_light.py +751 -0
  126. ldm/modules/image_degradation/utils_image.py +1024 -0
  127. ldm/modules/losses/__init__.py +1 -0
  128. ldm/modules/losses/contperceptual.py +175 -0
  129. ldm/modules/losses/vqperceptual.py +240 -0
  130. ldm/modules/x_transformer.py +752 -0
  131. ldm/simplet2i.py +13 -0
  132. ldm/util.py +264 -0
  133. scripts/configure_invokeai.py +797 -0
  134. scripts/create_root.py +43 -0
  135. scripts/dream.py +10 -0
  136. scripts/images2prompt.py +30 -0
  137. scripts/invoke.py +5 -0
  138. scripts/legacy_api.py +681 -0
  139. scripts/load_models.py +12 -0
  140. scripts/merge_embeddings.py +124 -0
  141. scripts/orig_scripts/img2img.py +296 -0
  142. scripts/orig_scripts/inpaint.py +98 -0
  143. scripts/orig_scripts/knn2img.py +397 -0
  144. scripts/orig_scripts/sample_diffusion.py +313 -0
  145. scripts/orig_scripts/train_searcher.py +147 -0
  146. scripts/orig_scripts/txt2img.py +324 -0
  147. scripts/preload_models.py +12 -0
  148. scripts/sd-metadata.py +22 -0
  149. vendored/clip/clip/__init__.py +1 -0
  150. vendored/clip/clip/bpe_simple_vocab_16e6.txt.gz +0 -0
  151. vendored/clip/clip/clip.py +237 -0
  152. vendored/clip/clip/model.py +436 -0
  153. vendored/clip/clip/simple_tokenizer.py +132 -0
  154. vendored/clip/hubconf.py +42 -0
  155. vendored/clip/setup.py +21 -0
  156. vendored/clip/tests/test_consistency.py +25 -0
  157. vendored/clipseg/datasets/coco_wrapper.py +99 -0
  158. vendored/clipseg/datasets/pascal_zeroshot.py +60 -0
  159. vendored/clipseg/datasets/pfe_dataset.py +129 -0
  160. vendored/clipseg/datasets/phrasecut.py +380 -0
  161. vendored/clipseg/datasets/utils.py +68 -0
  162. vendored/clipseg/evaluation_utils.py +292 -0
  163. vendored/clipseg/general_utils.py +272 -0
  164. vendored/clipseg/metrics.py +271 -0
  165. vendored/clipseg/models/clipseg.py +538 -0
  166. vendored/clipseg/models/vitseg.py +286 -0
  167. vendored/clipseg/score.py +453 -0
  168. vendored/clipseg/setup.py +30 -0
  169. vendored/clipseg/training.py +266 -0
@@ -0,0 +1,751 @@
1
+ # -*- coding: utf-8 -*-
2
+ import numpy as np
3
+ import cv2
4
+ import torch
5
+
6
+ from functools import partial
7
+ import random
8
+ from scipy import ndimage
9
+ import scipy
10
+ import scipy.stats as ss
11
+ from scipy.interpolate import interp2d
12
+ from scipy.linalg import orth
13
+ import albumentations
14
+
15
+ import ldm.modules.image_degradation.utils_image as util
16
+
17
+ """
18
+ # --------------------------------------------
19
+ # Super-Resolution
20
+ # --------------------------------------------
21
+ #
22
+ # Kai Zhang (cskaizhang@gmail.com)
23
+ # https://github.com/cszn
24
+ # From 2019/03--2021/08
25
+ # --------------------------------------------
26
+ """
27
+
28
+
29
+ def modcrop_np(img, sf):
30
+ """
31
+ Args:
32
+ img: numpy image, WxH or WxHxC
33
+ sf: scale factor
34
+ Return:
35
+ cropped image
36
+ """
37
+ w, h = img.shape[:2]
38
+ im = np.copy(img)
39
+ return im[: w - w % sf, : h - h % sf, ...]
40
+
41
+
42
+ """
43
+ # --------------------------------------------
44
+ # anisotropic Gaussian kernels
45
+ # --------------------------------------------
46
+ """
47
+
48
+
49
+ def analytic_kernel(k):
50
+ """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
51
+ k_size = k.shape[0]
52
+ # Calculate the big kernels size
53
+ big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
54
+ # Loop over the small kernel to fill the big one
55
+ for r in range(k_size):
56
+ for c in range(k_size):
57
+ big_k[2 * r : 2 * r + k_size, 2 * c : 2 * c + k_size] += (
58
+ k[r, c] * k
59
+ )
60
+ # Crop the edges of the big kernel to ignore very small values and increase run time of SR
61
+ crop = k_size // 2
62
+ cropped_big_k = big_k[crop:-crop, crop:-crop]
63
+ # Normalize to 1
64
+ return cropped_big_k / cropped_big_k.sum()
65
+
66
+
67
+ def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
68
+ """generate an anisotropic Gaussian kernel
69
+ Args:
70
+ ksize : e.g., 15, kernel size
71
+ theta : [0, pi], rotation angle range
72
+ l1 : [0.1,50], scaling of eigenvalues
73
+ l2 : [0.1,l1], scaling of eigenvalues
74
+ If l1 = l2, will get an isotropic Gaussian kernel.
75
+ Returns:
76
+ k : kernel
77
+ """
78
+
79
+ v = np.dot(
80
+ np.array(
81
+ [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]
82
+ ),
83
+ np.array([1.0, 0.0]),
84
+ )
85
+ V = np.array([[v[0], v[1]], [v[1], -v[0]]])
86
+ D = np.array([[l1, 0], [0, l2]])
87
+ Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
88
+ k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
89
+
90
+ return k
91
+
92
+
93
+ def gm_blur_kernel(mean, cov, size=15):
94
+ center = size / 2.0 + 0.5
95
+ k = np.zeros([size, size])
96
+ for y in range(size):
97
+ for x in range(size):
98
+ cy = y - center + 1
99
+ cx = x - center + 1
100
+ k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
101
+
102
+ k = k / np.sum(k)
103
+ return k
104
+
105
+
106
+ def shift_pixel(x, sf, upper_left=True):
107
+ """shift pixel for super-resolution with different scale factors
108
+ Args:
109
+ x: WxHxC or WxH
110
+ sf: scale factor
111
+ upper_left: shift direction
112
+ """
113
+ h, w = x.shape[:2]
114
+ shift = (sf - 1) * 0.5
115
+ xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
116
+ if upper_left:
117
+ x1 = xv + shift
118
+ y1 = yv + shift
119
+ else:
120
+ x1 = xv - shift
121
+ y1 = yv - shift
122
+
123
+ x1 = np.clip(x1, 0, w - 1)
124
+ y1 = np.clip(y1, 0, h - 1)
125
+
126
+ if x.ndim == 2:
127
+ x = interp2d(xv, yv, x)(x1, y1)
128
+ if x.ndim == 3:
129
+ for i in range(x.shape[-1]):
130
+ x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
131
+
132
+ return x
133
+
134
+
135
+ def blur(x, k):
136
+ """
137
+ x: image, NxcxHxW
138
+ k: kernel, Nx1xhxw
139
+ """
140
+ n, c = x.shape[:2]
141
+ p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2
142
+ x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
143
+ k = k.repeat(1, c, 1, 1)
144
+ k = k.view(-1, 1, k.shape[2], k.shape[3])
145
+ x = x.view(1, -1, x.shape[2], x.shape[3])
146
+ x = torch.nn.functional.conv2d(
147
+ x, k, bias=None, stride=1, padding=0, groups=n * c
148
+ )
149
+ x = x.view(n, c, x.shape[2], x.shape[3])
150
+
151
+ return x
152
+
153
+
154
+ def gen_kernel(
155
+ k_size=np.array([15, 15]),
156
+ scale_factor=np.array([4, 4]),
157
+ min_var=0.6,
158
+ max_var=10.0,
159
+ noise_level=0,
160
+ ):
161
+ """ "
162
+ # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
163
+ # Kai Zhang
164
+ # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var
165
+ # max_var = 2.5 * sf
166
+ """
167
+ # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
168
+ lambda_1 = min_var + np.random.rand() * (max_var - min_var)
169
+ lambda_2 = min_var + np.random.rand() * (max_var - min_var)
170
+ theta = np.random.rand() * np.pi # random theta
171
+ noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
172
+
173
+ # Set COV matrix using Lambdas and Theta
174
+ LAMBDA = np.diag([lambda_1, lambda_2])
175
+ Q = np.array(
176
+ [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]
177
+ )
178
+ SIGMA = Q @ LAMBDA @ Q.T
179
+ INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
180
+
181
+ # Set expectation position (shifting kernel for aligned image)
182
+ MU = k_size // 2 - 0.5 * (
183
+ scale_factor - 1
184
+ ) # - 0.5 * (scale_factor - k_size % 2)
185
+ MU = MU[None, None, :, None]
186
+
187
+ # Create meshgrid for Gaussian
188
+ [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
189
+ Z = np.stack([X, Y], 2)[:, :, :, None]
190
+
191
+ # Calcualte Gaussian for every pixel of the kernel
192
+ ZZ = Z - MU
193
+ ZZ_t = ZZ.transpose(0, 1, 3, 2)
194
+ raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
195
+
196
+ # shift the kernel so it will be centered
197
+ # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
198
+
199
+ # Normalize the kernel and return
200
+ # kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
201
+ kernel = raw_kernel / np.sum(raw_kernel)
202
+ return kernel
203
+
204
+
205
+ def fspecial_gaussian(hsize, sigma):
206
+ hsize = [hsize, hsize]
207
+ siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0]
208
+ std = sigma
209
+ [x, y] = np.meshgrid(
210
+ np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1)
211
+ )
212
+ arg = -(x * x + y * y) / (2 * std * std)
213
+ h = np.exp(arg)
214
+ h[h < scipy.finfo(float).eps * h.max()] = 0
215
+ sumh = h.sum()
216
+ if sumh != 0:
217
+ h = h / sumh
218
+ return h
219
+
220
+
221
+ def fspecial_laplacian(alpha):
222
+ alpha = max([0, min([alpha, 1])])
223
+ h1 = alpha / (alpha + 1)
224
+ h2 = (1 - alpha) / (alpha + 1)
225
+ h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]]
226
+ h = np.array(h)
227
+ return h
228
+
229
+
230
+ def fspecial(filter_type, *args, **kwargs):
231
+ """
232
+ python code from:
233
+ https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
234
+ """
235
+ if filter_type == 'gaussian':
236
+ return fspecial_gaussian(*args, **kwargs)
237
+ if filter_type == 'laplacian':
238
+ return fspecial_laplacian(*args, **kwargs)
239
+
240
+
241
+ """
242
+ # --------------------------------------------
243
+ # degradation models
244
+ # --------------------------------------------
245
+ """
246
+
247
+
248
+ def bicubic_degradation(x, sf=3):
249
+ """
250
+ Args:
251
+ x: HxWxC image, [0, 1]
252
+ sf: down-scale factor
253
+ Return:
254
+ bicubicly downsampled LR image
255
+ """
256
+ x = util.imresize_np(x, scale=1 / sf)
257
+ return x
258
+
259
+
260
+ def srmd_degradation(x, k, sf=3):
261
+ """blur + bicubic downsampling
262
+ Args:
263
+ x: HxWxC image, [0, 1]
264
+ k: hxw, double
265
+ sf: down-scale factor
266
+ Return:
267
+ downsampled LR image
268
+ Reference:
269
+ @inproceedings{zhang2018learning,
270
+ title={Learning a single convolutional super-resolution network for multiple degradations},
271
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
272
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
273
+ pages={3262--3271},
274
+ year={2018}
275
+ }
276
+ """
277
+ x = ndimage.filters.convolve(
278
+ x, np.expand_dims(k, axis=2), mode='wrap'
279
+ ) # 'nearest' | 'mirror'
280
+ x = bicubic_degradation(x, sf=sf)
281
+ return x
282
+
283
+
284
+ def dpsr_degradation(x, k, sf=3):
285
+ """bicubic downsampling + blur
286
+ Args:
287
+ x: HxWxC image, [0, 1]
288
+ k: hxw, double
289
+ sf: down-scale factor
290
+ Return:
291
+ downsampled LR image
292
+ Reference:
293
+ @inproceedings{zhang2019deep,
294
+ title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
295
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
296
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
297
+ pages={1671--1681},
298
+ year={2019}
299
+ }
300
+ """
301
+ x = bicubic_degradation(x, sf=sf)
302
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
303
+ return x
304
+
305
+
306
+ def classical_degradation(x, k, sf=3):
307
+ """blur + downsampling
308
+ Args:
309
+ x: HxWxC image, [0, 1]/[0, 255]
310
+ k: hxw, double
311
+ sf: down-scale factor
312
+ Return:
313
+ downsampled LR image
314
+ """
315
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
316
+ # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
317
+ st = 0
318
+ return x[st::sf, st::sf, ...]
319
+
320
+
321
+ def add_sharpening(img, weight=0.5, radius=50, threshold=10):
322
+ """USM sharpening. borrowed from real-ESRGAN
323
+ Input image: I; Blurry image: B.
324
+ 1. K = I + weight * (I - B)
325
+ 2. Mask = 1 if abs(I - B) > threshold, else: 0
326
+ 3. Blur mask:
327
+ 4. Out = Mask * K + (1 - Mask) * I
328
+ Args:
329
+ img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
330
+ weight (float): Sharp weight. Default: 1.
331
+ radius (float): Kernel size of Gaussian blur. Default: 50.
332
+ threshold (int):
333
+ """
334
+ if radius % 2 == 0:
335
+ radius += 1
336
+ blur = cv2.GaussianBlur(img, (radius, radius), 0)
337
+ residual = img - blur
338
+ mask = np.abs(residual) * 255 > threshold
339
+ mask = mask.astype('float32')
340
+ soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
341
+
342
+ K = img + weight * residual
343
+ K = np.clip(K, 0, 1)
344
+ return soft_mask * K + (1 - soft_mask) * img
345
+
346
+
347
+ def add_blur(img, sf=4):
348
+ wd2 = 4.0 + sf
349
+ wd = 2.0 + 0.2 * sf
350
+
351
+ wd2 = wd2 / 4
352
+ wd = wd / 4
353
+
354
+ if random.random() < 0.5:
355
+ l1 = wd2 * random.random()
356
+ l2 = wd2 * random.random()
357
+ k = anisotropic_Gaussian(
358
+ ksize=random.randint(2, 11) + 3,
359
+ theta=random.random() * np.pi,
360
+ l1=l1,
361
+ l2=l2,
362
+ )
363
+ else:
364
+ k = fspecial(
365
+ 'gaussian', random.randint(2, 4) + 3, wd * random.random()
366
+ )
367
+ img = ndimage.filters.convolve(
368
+ img, np.expand_dims(k, axis=2), mode='mirror'
369
+ )
370
+
371
+ return img
372
+
373
+
374
+ def add_resize(img, sf=4):
375
+ rnum = np.random.rand()
376
+ if rnum > 0.8: # up
377
+ sf1 = random.uniform(1, 2)
378
+ elif rnum < 0.7: # down
379
+ sf1 = random.uniform(0.5 / sf, 1)
380
+ else:
381
+ sf1 = 1.0
382
+ img = cv2.resize(
383
+ img,
384
+ (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])),
385
+ interpolation=random.choice([1, 2, 3]),
386
+ )
387
+ img = np.clip(img, 0.0, 1.0)
388
+
389
+ return img
390
+
391
+
392
+ # def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
393
+ # noise_level = random.randint(noise_level1, noise_level2)
394
+ # rnum = np.random.rand()
395
+ # if rnum > 0.6: # add color Gaussian noise
396
+ # img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
397
+ # elif rnum < 0.4: # add grayscale Gaussian noise
398
+ # img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
399
+ # else: # add noise
400
+ # L = noise_level2 / 255.
401
+ # D = np.diag(np.random.rand(3))
402
+ # U = orth(np.random.rand(3, 3))
403
+ # conv = np.dot(np.dot(np.transpose(U), D), U)
404
+ # img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
405
+ # img = np.clip(img, 0.0, 1.0)
406
+ # return img
407
+
408
+
409
+ def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
410
+ noise_level = random.randint(noise_level1, noise_level2)
411
+ rnum = np.random.rand()
412
+ if rnum > 0.6: # add color Gaussian noise
413
+ img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(
414
+ np.float32
415
+ )
416
+ elif rnum < 0.4: # add grayscale Gaussian noise
417
+ img = img + np.random.normal(
418
+ 0, noise_level / 255.0, (*img.shape[:2], 1)
419
+ ).astype(np.float32)
420
+ else: # add noise
421
+ L = noise_level2 / 255.0
422
+ D = np.diag(np.random.rand(3))
423
+ U = orth(np.random.rand(3, 3))
424
+ conv = np.dot(np.dot(np.transpose(U), D), U)
425
+ img = img + np.random.multivariate_normal(
426
+ [0, 0, 0], np.abs(L**2 * conv), img.shape[:2]
427
+ ).astype(np.float32)
428
+ img = np.clip(img, 0.0, 1.0)
429
+ return img
430
+
431
+
432
+ def add_speckle_noise(img, noise_level1=2, noise_level2=25):
433
+ noise_level = random.randint(noise_level1, noise_level2)
434
+ img = np.clip(img, 0.0, 1.0)
435
+ rnum = random.random()
436
+ if rnum > 0.6:
437
+ img += img * np.random.normal(
438
+ 0, noise_level / 255.0, img.shape
439
+ ).astype(np.float32)
440
+ elif rnum < 0.4:
441
+ img += img * np.random.normal(
442
+ 0, noise_level / 255.0, (*img.shape[:2], 1)
443
+ ).astype(np.float32)
444
+ else:
445
+ L = noise_level2 / 255.0
446
+ D = np.diag(np.random.rand(3))
447
+ U = orth(np.random.rand(3, 3))
448
+ conv = np.dot(np.dot(np.transpose(U), D), U)
449
+ img += img * np.random.multivariate_normal(
450
+ [0, 0, 0], np.abs(L**2 * conv), img.shape[:2]
451
+ ).astype(np.float32)
452
+ img = np.clip(img, 0.0, 1.0)
453
+ return img
454
+
455
+
456
+ def add_Poisson_noise(img):
457
+ img = np.clip((img * 255.0).round(), 0, 255) / 255.0
458
+ vals = 10 ** (2 * random.random() + 2.0) # [2, 4]
459
+ if random.random() < 0.5:
460
+ img = np.random.poisson(img * vals).astype(np.float32) / vals
461
+ else:
462
+ img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114])
463
+ img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.0
464
+ noise_gray = (
465
+ np.random.poisson(img_gray * vals).astype(np.float32) / vals
466
+ - img_gray
467
+ )
468
+ img += noise_gray[:, :, np.newaxis]
469
+ img = np.clip(img, 0.0, 1.0)
470
+ return img
471
+
472
+
473
+ def add_JPEG_noise(img):
474
+ quality_factor = random.randint(80, 95)
475
+ img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
476
+ result, encimg = cv2.imencode(
477
+ '.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor]
478
+ )
479
+ img = cv2.imdecode(encimg, 1)
480
+ img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
481
+ return img
482
+
483
+
484
+ def random_crop(lq, hq, sf=4, lq_patchsize=64):
485
+ h, w = lq.shape[:2]
486
+ rnd_h = random.randint(0, h - lq_patchsize)
487
+ rnd_w = random.randint(0, w - lq_patchsize)
488
+ lq = lq[rnd_h : rnd_h + lq_patchsize, rnd_w : rnd_w + lq_patchsize, :]
489
+
490
+ rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
491
+ hq = hq[
492
+ rnd_h_H : rnd_h_H + lq_patchsize * sf,
493
+ rnd_w_H : rnd_w_H + lq_patchsize * sf,
494
+ :,
495
+ ]
496
+ return lq, hq
497
+
498
+
499
+ def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None):
500
+ """
501
+ This is the degradation model of BSRGAN from the paper
502
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
503
+ ----------
504
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
505
+ sf: scale factor
506
+ isp_model: camera ISP model
507
+ Returns
508
+ -------
509
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
510
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
511
+ """
512
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
513
+ sf_ori = sf
514
+
515
+ h1, w1 = img.shape[:2]
516
+ img = img.copy()[: w1 - w1 % sf, : h1 - h1 % sf, ...] # mod crop
517
+ h, w = img.shape[:2]
518
+
519
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
520
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
521
+
522
+ hq = img.copy()
523
+
524
+ if sf == 4 and random.random() < scale2_prob: # downsample1
525
+ if np.random.rand() < 0.5:
526
+ img = cv2.resize(
527
+ img,
528
+ (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])),
529
+ interpolation=random.choice([1, 2, 3]),
530
+ )
531
+ else:
532
+ img = util.imresize_np(img, 1 / 2, True)
533
+ img = np.clip(img, 0.0, 1.0)
534
+ sf = 2
535
+
536
+ shuffle_order = random.sample(range(7), 7)
537
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
538
+ if idx1 > idx2: # keep downsample3 last
539
+ shuffle_order[idx1], shuffle_order[idx2] = (
540
+ shuffle_order[idx2],
541
+ shuffle_order[idx1],
542
+ )
543
+
544
+ for i in shuffle_order:
545
+
546
+ if i == 0:
547
+ img = add_blur(img, sf=sf)
548
+
549
+ elif i == 1:
550
+ img = add_blur(img, sf=sf)
551
+
552
+ elif i == 2:
553
+ a, b = img.shape[1], img.shape[0]
554
+ # downsample2
555
+ if random.random() < 0.75:
556
+ sf1 = random.uniform(1, 2 * sf)
557
+ img = cv2.resize(
558
+ img,
559
+ (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])),
560
+ interpolation=random.choice([1, 2, 3]),
561
+ )
562
+ else:
563
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
564
+ k_shifted = shift_pixel(k, sf)
565
+ k_shifted = (
566
+ k_shifted / k_shifted.sum()
567
+ ) # blur with shifted kernel
568
+ img = ndimage.filters.convolve(
569
+ img, np.expand_dims(k_shifted, axis=2), mode='mirror'
570
+ )
571
+ img = img[0::sf, 0::sf, ...] # nearest downsampling
572
+ img = np.clip(img, 0.0, 1.0)
573
+
574
+ elif i == 3:
575
+ # downsample3
576
+ img = cv2.resize(
577
+ img,
578
+ (int(1 / sf * a), int(1 / sf * b)),
579
+ interpolation=random.choice([1, 2, 3]),
580
+ )
581
+ img = np.clip(img, 0.0, 1.0)
582
+
583
+ elif i == 4:
584
+ # add Gaussian noise
585
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=8)
586
+
587
+ elif i == 5:
588
+ # add JPEG noise
589
+ if random.random() < jpeg_prob:
590
+ img = add_JPEG_noise(img)
591
+
592
+ elif i == 6:
593
+ # add processed camera sensor noise
594
+ if random.random() < isp_prob and isp_model is not None:
595
+ with torch.no_grad():
596
+ img, hq = isp_model.forward(img.copy(), hq)
597
+
598
+ # add final JPEG compression noise
599
+ img = add_JPEG_noise(img)
600
+
601
+ # random crop
602
+ img, hq = random_crop(img, hq, sf_ori, lq_patchsize)
603
+
604
+ return img, hq
605
+
606
+
607
+ # todo no isp_model?
608
+ def degradation_bsrgan_variant(image, sf=4, isp_model=None):
609
+ """
610
+ This is the degradation model of BSRGAN from the paper
611
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
612
+ ----------
613
+ sf: scale factor
614
+ isp_model: camera ISP model
615
+ Returns
616
+ -------
617
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
618
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
619
+ """
620
+ image = util.uint2single(image)
621
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
622
+ sf_ori = sf
623
+
624
+ h1, w1 = image.shape[:2]
625
+ image = image.copy()[: w1 - w1 % sf, : h1 - h1 % sf, ...] # mod crop
626
+ h, w = image.shape[:2]
627
+
628
+ hq = image.copy()
629
+
630
+ if sf == 4 and random.random() < scale2_prob: # downsample1
631
+ if np.random.rand() < 0.5:
632
+ image = cv2.resize(
633
+ image,
634
+ (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
635
+ interpolation=random.choice([1, 2, 3]),
636
+ )
637
+ else:
638
+ image = util.imresize_np(image, 1 / 2, True)
639
+ image = np.clip(image, 0.0, 1.0)
640
+ sf = 2
641
+
642
+ shuffle_order = random.sample(range(7), 7)
643
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
644
+ if idx1 > idx2: # keep downsample3 last
645
+ shuffle_order[idx1], shuffle_order[idx2] = (
646
+ shuffle_order[idx2],
647
+ shuffle_order[idx1],
648
+ )
649
+
650
+ for i in shuffle_order:
651
+
652
+ if i == 0:
653
+ image = add_blur(image, sf=sf)
654
+
655
+ # elif i == 1:
656
+ # image = add_blur(image, sf=sf)
657
+
658
+ if i == 0:
659
+ pass
660
+
661
+ elif i == 2:
662
+ a, b = image.shape[1], image.shape[0]
663
+ # downsample2
664
+ if random.random() < 0.8:
665
+ sf1 = random.uniform(1, 2 * sf)
666
+ image = cv2.resize(
667
+ image,
668
+ (
669
+ int(1 / sf1 * image.shape[1]),
670
+ int(1 / sf1 * image.shape[0]),
671
+ ),
672
+ interpolation=random.choice([1, 2, 3]),
673
+ )
674
+ else:
675
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
676
+ k_shifted = shift_pixel(k, sf)
677
+ k_shifted = (
678
+ k_shifted / k_shifted.sum()
679
+ ) # blur with shifted kernel
680
+ image = ndimage.filters.convolve(
681
+ image, np.expand_dims(k_shifted, axis=2), mode='mirror'
682
+ )
683
+ image = image[0::sf, 0::sf, ...] # nearest downsampling
684
+
685
+ image = np.clip(image, 0.0, 1.0)
686
+
687
+ elif i == 3:
688
+ # downsample3
689
+ image = cv2.resize(
690
+ image,
691
+ (int(1 / sf * a), int(1 / sf * b)),
692
+ interpolation=random.choice([1, 2, 3]),
693
+ )
694
+ image = np.clip(image, 0.0, 1.0)
695
+
696
+ elif i == 4:
697
+ # add Gaussian noise
698
+ image = add_Gaussian_noise(image, noise_level1=1, noise_level2=2)
699
+
700
+ elif i == 5:
701
+ # add JPEG noise
702
+ if random.random() < jpeg_prob:
703
+ image = add_JPEG_noise(image)
704
+ #
705
+ # elif i == 6:
706
+ # # add processed camera sensor noise
707
+ # if random.random() < isp_prob and isp_model is not None:
708
+ # with torch.no_grad():
709
+ # img, hq = isp_model.forward(img.copy(), hq)
710
+
711
+ # add final JPEG compression noise
712
+ image = add_JPEG_noise(image)
713
+ image = util.single2uint(image)
714
+ example = {'image': image}
715
+ return example
716
+
717
+
718
+ if __name__ == '__main__':
719
+ print('hey')
720
+ img = util.imread_uint('utils/test.png', 3)
721
+ img = img[:448, :448]
722
+ h = img.shape[0] // 4
723
+ print('resizing to', h)
724
+ sf = 4
725
+ deg_fn = partial(degradation_bsrgan_variant, sf=sf)
726
+ for i in range(20):
727
+ print(i)
728
+ img_hq = img
729
+ img_lq = deg_fn(img)['image']
730
+ img_hq, img_lq = util.uint2single(img_hq), util.uint2single(img_lq)
731
+ print(img_lq)
732
+ img_lq_bicubic = albumentations.SmallestMaxSize(
733
+ max_size=h, interpolation=cv2.INTER_CUBIC
734
+ )(image=img_hq)['image']
735
+ print(img_lq.shape)
736
+ print('bicubic', img_lq_bicubic.shape)
737
+ print(img_hq.shape)
738
+ lq_nearest = cv2.resize(
739
+ util.single2uint(img_lq),
740
+ (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
741
+ interpolation=0,
742
+ )
743
+ lq_bicubic_nearest = cv2.resize(
744
+ util.single2uint(img_lq_bicubic),
745
+ (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
746
+ interpolation=0,
747
+ )
748
+ img_concat = np.concatenate(
749
+ [lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1
750
+ )
751
+ util.imsave(img_concat, str(i) + '.png')