Dynamatrix/DiffBIR-OpenXLab
0
1# https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/data/degradations.py2import cv23import math4import numpy as np5import random6import torch7from scipy import special8from scipy.stats import multivariate_normal9from torchvision.transforms.functional_tensor import rgb_to_grayscale10 11# -------------------------------------------------------------------- #12# --------------------------- blur kernels --------------------------- #13# -------------------------------------------------------------------- #14 15 16# --------------------------- util functions --------------------------- #17def sigma_matrix2(sig_x, sig_y, theta):18 """Calculate the rotated sigma matrix (two dimensional matrix).19 20 Args:21 sig_x (float):22 sig_y (float):23 theta (float): Radian measurement.24 25 Returns:26 ndarray: Rotated sigma matrix.27 """28 d_matrix = np.array([[sig_x**2, 0], [0, sig_y**2]])29 u_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])30 return np.dot(u_matrix, np.dot(d_matrix, u_matrix.T))31 32 33def mesh_grid(kernel_size):34 """Generate the mesh grid, centering at zero.35 36 Args:37 kernel_size (int):38 39 Returns:40 xy (ndarray): with the shape (kernel_size, kernel_size, 2)41 xx (ndarray): with the shape (kernel_size, kernel_size)42 yy (ndarray): with the shape (kernel_size, kernel_size)43 """44 ax = np.arange(-kernel_size // 2 + 1., kernel_size // 2 + 1.)45 xx, yy = np.meshgrid(ax, ax)46 xy = np.hstack((xx.reshape((kernel_size * kernel_size, 1)), yy.reshape(kernel_size * kernel_size,47 1))).reshape(kernel_size, kernel_size, 2)48 return xy, xx, yy49 50 51def pdf2(sigma_matrix, grid):52 """Calculate PDF of the bivariate Gaussian distribution.53 54 Args:55 sigma_matrix (ndarray): with the shape (2, 2)56 grid (ndarray): generated by :func:`mesh_grid`,57 with the shape (K, K, 2), K is the kernel size.58 59 Returns:60 kernel (ndarrray): un-normalized kernel.61 """62 inverse_sigma = np.linalg.inv(sigma_matrix)63 kernel = np.exp(-0.5 * np.sum(np.dot(grid, inverse_sigma) * grid, 2))64 return kernel65 66 67def cdf2(d_matrix, grid):68 """Calculate the CDF of the standard bivariate Gaussian distribution.69 Used in skewed Gaussian distribution.70 71 Args:72 d_matrix (ndarrasy): skew matrix.73 grid (ndarray): generated by :func:`mesh_grid`,74 with the shape (K, K, 2), K is the kernel size.75 76 Returns:77 cdf (ndarray): skewed cdf.78 """79 rv = multivariate_normal([0, 0], [[1, 0], [0, 1]])80 grid = np.dot(grid, d_matrix)81 cdf = rv.cdf(grid)82 return cdf83 84 85def bivariate_Gaussian(kernel_size, sig_x, sig_y, theta, grid=None, isotropic=True):86 """Generate a bivariate isotropic or anisotropic Gaussian kernel.87 88 In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.89 90 Args:91 kernel_size (int):92 sig_x (float):93 sig_y (float):94 theta (float): Radian measurement.95 grid (ndarray, optional): generated by :func:`mesh_grid`,96 with the shape (K, K, 2), K is the kernel size. Default: None97 isotropic (bool):98 99 Returns:100 kernel (ndarray): normalized kernel.101 """102 if grid is None:103 grid, _, _ = mesh_grid(kernel_size)104 if isotropic:105 sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])106 else:107 sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)108 kernel = pdf2(sigma_matrix, grid)109 kernel = kernel / np.sum(kernel)110 return kernel111 112 113def bivariate_generalized_Gaussian(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):114 """Generate a bivariate generalized Gaussian kernel.115 116 ``Paper: Parameter Estimation For Multivariate Generalized Gaussian Distributions``117 118 In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.119 120 Args:121 kernel_size (int):122 sig_x (float):123 sig_y (float):124 theta (float): Radian measurement.125 beta (float): shape parameter, beta = 1 is the normal distribution.126 grid (ndarray, optional): generated by :func:`mesh_grid`,127 with the shape (K, K, 2), K is the kernel size. Default: None128 129 Returns:130 kernel (ndarray): normalized kernel.131 """132 if grid is None:133 grid, _, _ = mesh_grid(kernel_size)134 if isotropic:135 sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])136 else:137 sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)138 inverse_sigma = np.linalg.inv(sigma_matrix)139 kernel = np.exp(-0.5 * np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta))140 kernel = kernel / np.sum(kernel)141 return kernel142 143 144def bivariate_plateau(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):145 """Generate a plateau-like anisotropic kernel.146 147 1 / (1+x^(beta))148 149 Reference: https://stats.stackexchange.com/questions/203629/is-there-a-plateau-shaped-distribution150 151 In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.152 153 Args:154 kernel_size (int):155 sig_x (float):156 sig_y (float):157 theta (float): Radian measurement.158 beta (float): shape parameter, beta = 1 is the normal distribution.159 grid (ndarray, optional): generated by :func:`mesh_grid`,160 with the shape (K, K, 2), K is the kernel size. Default: None161 162 Returns:163 kernel (ndarray): normalized kernel.164 """165 if grid is None:166 grid, _, _ = mesh_grid(kernel_size)167 if isotropic:168 sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])169 else:170 sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)171 inverse_sigma = np.linalg.inv(sigma_matrix)172 kernel = np.reciprocal(np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta) + 1)173 kernel = kernel / np.sum(kernel)174 return kernel175 176 177def random_bivariate_Gaussian(kernel_size,178 sigma_x_range,179 sigma_y_range,180 rotation_range,181 noise_range=None,182 isotropic=True):183 """Randomly generate bivariate isotropic or anisotropic Gaussian kernels.184 185 In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.186 187 Args:188 kernel_size (int):189 sigma_x_range (tuple): [0.6, 5]190 sigma_y_range (tuple): [0.6, 5]191 rotation range (tuple): [-math.pi, math.pi]192 noise_range(tuple, optional): multiplicative kernel noise,193 [0.75, 1.25]. Default: None194 195 Returns:196 kernel (ndarray):197 """198 assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'199 assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'200 sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])201 if isotropic is False:202 assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'203 assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'204 sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])205 rotation = np.random.uniform(rotation_range[0], rotation_range[1])206 else:207 sigma_y = sigma_x208 rotation = 0209 210 kernel = bivariate_Gaussian(kernel_size, sigma_x, sigma_y, rotation, isotropic=isotropic)211 212 # add multiplicative noise213 if noise_range is not None:214 assert noise_range[0] < noise_range[1], 'Wrong noise range.'215 noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)216 kernel = kernel * noise217 kernel = kernel / np.sum(kernel)218 return kernel219 220 221def random_bivariate_generalized_Gaussian(kernel_size,222 sigma_x_range,223 sigma_y_range,224 rotation_range,225 beta_range,226 noise_range=None,227 isotropic=True):228 """Randomly generate bivariate generalized Gaussian kernels.229 230 In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.231 232 Args:233 kernel_size (int):234 sigma_x_range (tuple): [0.6, 5]235 sigma_y_range (tuple): [0.6, 5]236 rotation range (tuple): [-math.pi, math.pi]237 beta_range (tuple): [0.5, 8]238 noise_range(tuple, optional): multiplicative kernel noise,239 [0.75, 1.25]. Default: None240 241 Returns:242 kernel (ndarray):243 """244 assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'245 assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'246 sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])247 if isotropic is False:248 assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'249 assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'250 sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])251 rotation = np.random.uniform(rotation_range[0], rotation_range[1])252 else:253 sigma_y = sigma_x254 rotation = 0255 256 # assume beta_range[0] < 1 < beta_range[1]257 if np.random.uniform() < 0.5:258 beta = np.random.uniform(beta_range[0], 1)259 else:260 beta = np.random.uniform(1, beta_range[1])261 262 kernel = bivariate_generalized_Gaussian(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)263 264 # add multiplicative noise265 if noise_range is not None:266 assert noise_range[0] < noise_range[1], 'Wrong noise range.'267 noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)268 kernel = kernel * noise269 kernel = kernel / np.sum(kernel)270 return kernel271 272 273def random_bivariate_plateau(kernel_size,274 sigma_x_range,275 sigma_y_range,276 rotation_range,277 beta_range,278 noise_range=None,279 isotropic=True):280 """Randomly generate bivariate plateau kernels.281 282 In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.283 284 Args:285 kernel_size (int):286 sigma_x_range (tuple): [0.6, 5]287 sigma_y_range (tuple): [0.6, 5]288 rotation range (tuple): [-math.pi/2, math.pi/2]289 beta_range (tuple): [1, 4]290 noise_range(tuple, optional): multiplicative kernel noise,291 [0.75, 1.25]. Default: None292 293 Returns:294 kernel (ndarray):295 """296 assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'297 assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'298 sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])299 if isotropic is False:300 assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'301 assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'302 sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])303 rotation = np.random.uniform(rotation_range[0], rotation_range[1])304 else:305 sigma_y = sigma_x306 rotation = 0307 308 # TODO: this may be not proper309 if np.random.uniform() < 0.5:310 beta = np.random.uniform(beta_range[0], 1)311 else:312 beta = np.random.uniform(1, beta_range[1])313 314 kernel = bivariate_plateau(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)315 # add multiplicative noise316 if noise_range is not None:317 assert noise_range[0] < noise_range[1], 'Wrong noise range.'318 noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)319 kernel = kernel * noise320 kernel = kernel / np.sum(kernel)321 322 return kernel323 324 325def random_mixed_kernels(kernel_list,326 kernel_prob,327 kernel_size=21,328 sigma_x_range=(0.6, 5),329 sigma_y_range=(0.6, 5),330 rotation_range=(-math.pi, math.pi),331 betag_range=(0.5, 8),332 betap_range=(0.5, 8),333 noise_range=None):334 """Randomly generate mixed kernels.335 336 Args:337 kernel_list (tuple): a list name of kernel types,338 support ['iso', 'aniso', 'skew', 'generalized', 'plateau_iso',339 'plateau_aniso']340 kernel_prob (tuple): corresponding kernel probability for each341 kernel type342 kernel_size (int):343 sigma_x_range (tuple): [0.6, 5]344 sigma_y_range (tuple): [0.6, 5]345 rotation range (tuple): [-math.pi, math.pi]346 beta_range (tuple): [0.5, 8]347 noise_range(tuple, optional): multiplicative kernel noise,348 [0.75, 1.25]. Default: None349 350 Returns:351 kernel (ndarray):352 """353 kernel_type = random.choices(kernel_list, kernel_prob)[0]354 if kernel_type == 'iso':355 kernel = random_bivariate_Gaussian(356 kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=True)357 elif kernel_type == 'aniso':358 kernel = random_bivariate_Gaussian(359 kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=False)360 elif kernel_type == 'generalized_iso':361 kernel = random_bivariate_generalized_Gaussian(362 kernel_size,363 sigma_x_range,364 sigma_y_range,365 rotation_range,366 betag_range,367 noise_range=noise_range,368 isotropic=True)369 elif kernel_type == 'generalized_aniso':370 kernel = random_bivariate_generalized_Gaussian(371 kernel_size,372 sigma_x_range,373 sigma_y_range,374 rotation_range,375 betag_range,376 noise_range=noise_range,377 isotropic=False)378 elif kernel_type == 'plateau_iso':379 kernel = random_bivariate_plateau(380 kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=True)381 elif kernel_type == 'plateau_aniso':382 kernel = random_bivariate_plateau(383 kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=False)384 return kernel385 386 387np.seterr(divide='ignore', invalid='ignore')388 389 390def circular_lowpass_kernel(cutoff, kernel_size, pad_to=0):391 """2D sinc filter392 393 Reference: https://dsp.stackexchange.com/questions/58301/2-d-circularly-symmetric-low-pass-filter394 395 Args:396 cutoff (float): cutoff frequency in radians (pi is max)397 kernel_size (int): horizontal and vertical size, must be odd.398 pad_to (int): pad kernel size to desired size, must be odd or zero.399 """400 assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'401 kernel = np.fromfunction(402 lambda x, y: cutoff * special.j1(cutoff * np.sqrt(403 (x - (kernel_size - 1) / 2)**2 + (y - (kernel_size - 1) / 2)**2)) / (2 * np.pi * np.sqrt(404 (x - (kernel_size - 1) / 2)**2 + (y - (kernel_size - 1) / 2)**2)), [kernel_size, kernel_size])405 kernel[(kernel_size - 1) // 2, (kernel_size - 1) // 2] = cutoff**2 / (4 * np.pi)406 kernel = kernel / np.sum(kernel)407 if pad_to > kernel_size:408 pad_size = (pad_to - kernel_size) // 2409 kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))410 return kernel411 412 413# ------------------------------------------------------------- #414# --------------------------- noise --------------------------- #415# ------------------------------------------------------------- #416 417# ----------------------- Gaussian Noise ----------------------- #418 419 420def generate_gaussian_noise(img, sigma=10, gray_noise=False):421 """Generate Gaussian noise.422 423 Args:424 img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.425 sigma (float): Noise scale (measured in range 255). Default: 10.426 427 Returns:428 (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],429 float32.430 """431 if gray_noise:432 noise = np.float32(np.random.randn(*(img.shape[0:2]))) * sigma / 255.433 noise = np.expand_dims(noise, axis=2).repeat(3, axis=2)434 else:435 noise = np.float32(np.random.randn(*(img.shape))) * sigma / 255.436 return noise437 438 439def add_gaussian_noise(img, sigma=10, clip=True, rounds=False, gray_noise=False):440 """Add Gaussian noise.441 442 Args:443 img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.444 sigma (float): Noise scale (measured in range 255). Default: 10.445 446 Returns:447 (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],448 float32.449 """450 noise = generate_gaussian_noise(img, sigma, gray_noise)451 out = img + noise452 if clip and rounds:453 out = np.clip((out * 255.0).round(), 0, 255) / 255.454 elif clip:455 out = np.clip(out, 0, 1)456 elif rounds:457 out = (out * 255.0).round() / 255.458 return out459 460 461def generate_gaussian_noise_pt(img, sigma=10, gray_noise=0):462 """Add Gaussian noise (PyTorch version).463 464 Args:465 img (Tensor): Shape (b, c, h, w), range[0, 1], float32.466 scale (float | Tensor): Noise scale. Default: 1.0.467 468 Returns:469 (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],470 float32.471 """472 b, _, h, w = img.size()473 if not isinstance(sigma, (float, int)):474 sigma = sigma.view(img.size(0), 1, 1, 1)475 if isinstance(gray_noise, (float, int)):476 cal_gray_noise = gray_noise > 0477 else:478 gray_noise = gray_noise.view(b, 1, 1, 1)479 cal_gray_noise = torch.sum(gray_noise) > 0480 481 if cal_gray_noise:482 noise_gray = torch.randn(*img.size()[2:4], dtype=img.dtype, device=img.device) * sigma / 255.483 noise_gray = noise_gray.view(b, 1, h, w)484 485 # always calculate color noise486 noise = torch.randn(*img.size(), dtype=img.dtype, device=img.device) * sigma / 255.487 488 if cal_gray_noise:489 noise = noise * (1 - gray_noise) + noise_gray * gray_noise490 return noise491 492 493def add_gaussian_noise_pt(img, sigma=10, gray_noise=0, clip=True, rounds=False):494 """Add Gaussian noise (PyTorch version).495 496 Args:497 img (Tensor): Shape (b, c, h, w), range[0, 1], float32.498 scale (float | Tensor): Noise scale. Default: 1.0.499 500 Returns:501 (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],502 float32.503 """504 noise = generate_gaussian_noise_pt(img, sigma, gray_noise)505 out = img + noise506 if clip and rounds:507 out = torch.clamp((out * 255.0).round(), 0, 255) / 255.508 elif clip:509 out = torch.clamp(out, 0, 1)510 elif rounds:511 out = (out * 255.0).round() / 255.512 return out513 514 515# ----------------------- Random Gaussian Noise ----------------------- #516def random_generate_gaussian_noise(img, sigma_range=(0, 10), gray_prob=0):517 sigma = np.random.uniform(sigma_range[0], sigma_range[1])518 if np.random.uniform() < gray_prob:519 gray_noise = True520 else:521 gray_noise = False522 return generate_gaussian_noise(img, sigma, gray_noise)523 524 525def random_add_gaussian_noise(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):526 noise = random_generate_gaussian_noise(img, sigma_range, gray_prob)527 out = img + noise528 if clip and rounds:529 out = np.clip((out * 255.0).round(), 0, 255) / 255.530 elif clip:531 out = np.clip(out, 0, 1)532 elif rounds:533 out = (out * 255.0).round() / 255.534 return out535 536 537def random_generate_gaussian_noise_pt(img, sigma_range=(0, 10), gray_prob=0):538 sigma = torch.rand(539 img.size(0), dtype=img.dtype, device=img.device) * (sigma_range[1] - sigma_range[0]) + sigma_range[0]540 gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)541 gray_noise = (gray_noise < gray_prob).float()542 return generate_gaussian_noise_pt(img, sigma, gray_noise)543 544 545def random_add_gaussian_noise_pt(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):546 noise = random_generate_gaussian_noise_pt(img, sigma_range, gray_prob)547 out = img + noise548 if clip and rounds:549 out = torch.clamp((out * 255.0).round(), 0, 255) / 255.550 elif clip:551 out = torch.clamp(out, 0, 1)552 elif rounds:553 out = (out * 255.0).round() / 255.554 return out555 556 557# ----------------------- Poisson (Shot) Noise ----------------------- #558 559 560def generate_poisson_noise(img, scale=1.0, gray_noise=False):561 """Generate poisson noise.562 563 Reference: https://github.com/scikit-image/scikit-image/blob/main/skimage/util/noise.py#L37-L219564 565 Args:566 img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.567 scale (float): Noise scale. Default: 1.0.568 gray_noise (bool): Whether generate gray noise. Default: False.569 570 Returns:571 (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],572 float32.573 """574 if gray_noise:575 img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)576 # round and clip image for counting vals correctly577 img = np.clip((img * 255.0).round(), 0, 255) / 255.578 vals = len(np.unique(img))579 vals = 2**np.ceil(np.log2(vals))580 out = np.float32(np.random.poisson(img * vals) / float(vals))581 noise = out - img582 if gray_noise:583 noise = np.repeat(noise[:, :, np.newaxis], 3, axis=2)584 return noise * scale585 586 587def add_poisson_noise(img, scale=1.0, clip=True, rounds=False, gray_noise=False):588 """Add poisson noise.589 590 Args:591 img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.592 scale (float): Noise scale. Default: 1.0.593 gray_noise (bool): Whether generate gray noise. Default: False.594 595 Returns:596 (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],597 float32.598 """599 noise = generate_poisson_noise(img, scale, gray_noise)600 out = img + noise601 if clip and rounds:602 out = np.clip((out * 255.0).round(), 0, 255) / 255.603 elif clip:604 out = np.clip(out, 0, 1)605 elif rounds:606 out = (out * 255.0).round() / 255.607 return out608 609 610def generate_poisson_noise_pt(img, scale=1.0, gray_noise=0):611 """Generate a batch of poisson noise (PyTorch version)612 613 Args:614 img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.615 scale (float | Tensor): Noise scale. Number or Tensor with shape (b).616 Default: 1.0.617 gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).618 0 for False, 1 for True. Default: 0.619 620 Returns:621 (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],622 float32.623 """624 b, _, h, w = img.size()625 if isinstance(gray_noise, (float, int)):626 cal_gray_noise = gray_noise > 0627 else:628 gray_noise = gray_noise.view(b, 1, 1, 1)629 cal_gray_noise = torch.sum(gray_noise) > 0630 if cal_gray_noise:631 img_gray = rgb_to_grayscale(img, num_output_channels=1)632 # round and clip image for counting vals correctly633 img_gray = torch.clamp((img_gray * 255.0).round(), 0, 255) / 255.634 # use for-loop to get the unique values for each sample635 vals_list = [len(torch.unique(img_gray[i, :, :, :])) for i in range(b)]636 vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]637 vals = img_gray.new_tensor(vals_list).view(b, 1, 1, 1)638 out = torch.poisson(img_gray * vals) / vals639 noise_gray = out - img_gray640 noise_gray = noise_gray.expand(b, 3, h, w)641 642 # always calculate color noise643 # round and clip image for counting vals correctly644 img = torch.clamp((img * 255.0).round(), 0, 255) / 255.645 # use for-loop to get the unique values for each sample646 vals_list = [len(torch.unique(img[i, :, :, :])) for i in range(b)]647 vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]648 vals = img.new_tensor(vals_list).view(b, 1, 1, 1)649 out = torch.poisson(img * vals) / vals650 noise = out - img651 if cal_gray_noise:652 noise = noise * (1 - gray_noise) + noise_gray * gray_noise653 if not isinstance(scale, (float, int)):654 scale = scale.view(b, 1, 1, 1)655 return noise * scale656 657 658def add_poisson_noise_pt(img, scale=1.0, clip=True, rounds=False, gray_noise=0):659 """Add poisson noise to a batch of images (PyTorch version).660 661 Args:662 img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.663 scale (float | Tensor): Noise scale. Number or Tensor with shape (b).664 Default: 1.0.665 gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).666 0 for False, 1 for True. Default: 0.667 668 Returns:669 (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],670 float32.671 """672 noise = generate_poisson_noise_pt(img, scale, gray_noise)673 out = img + noise674 if clip and rounds:675 out = torch.clamp((out * 255.0).round(), 0, 255) / 255.676 elif clip:677 out = torch.clamp(out, 0, 1)678 elif rounds:679 out = (out * 255.0).round() / 255.680 return out681 682 683# ----------------------- Random Poisson (Shot) Noise ----------------------- #684 685 686def random_generate_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0):687 scale = np.random.uniform(scale_range[0], scale_range[1])688 if np.random.uniform() < gray_prob:689 gray_noise = True690 else:691 gray_noise = False692 return generate_poisson_noise(img, scale, gray_noise)693 694 695def random_add_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):696 noise = random_generate_poisson_noise(img, scale_range, gray_prob)697 out = img + noise698 if clip and rounds:699 out = np.clip((out * 255.0).round(), 0, 255) / 255.700 elif clip:701 out = np.clip(out, 0, 1)702 elif rounds:703 out = (out * 255.0).round() / 255.704 return out705 706 707def random_generate_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0):708 scale = torch.rand(709 img.size(0), dtype=img.dtype, device=img.device) * (scale_range[1] - scale_range[0]) + scale_range[0]710 gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)711 gray_noise = (gray_noise < gray_prob).float()712 return generate_poisson_noise_pt(img, scale, gray_noise)713 714 715def random_add_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):716 noise = random_generate_poisson_noise_pt(img, scale_range, gray_prob)717 out = img + noise718 if clip and rounds:719 out = torch.clamp((out * 255.0).round(), 0, 255) / 255.720 elif clip:721 out = torch.clamp(out, 0, 1)722 elif rounds:723 out = (out * 255.0).round() / 255.724 return out725 726 727# ------------------------------------------------------------------------ #728# --------------------------- JPEG compression --------------------------- #729# ------------------------------------------------------------------------ #730 731 732def add_jpg_compression(img, quality=90):733 """Add JPG compression artifacts.734 735 Args:736 img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.737 quality (float): JPG compression quality. 0 for lowest quality, 100 for738 best quality. Default: 90.739 740 Returns:741 (Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],742 float32.743 """744 img = np.clip(img, 0, 1)745 encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]746 _, encimg = cv2.imencode('.jpg', img * 255., encode_param)747 img = np.float32(cv2.imdecode(encimg, 1)) / 255.748 return img749 750 751def random_add_jpg_compression(img, quality_range=(90, 100)):752 """Randomly add JPG compression artifacts.753 754 Args:755 img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.756 quality_range (tuple[float] | list[float]): JPG compression quality757 range. 0 for lowest quality, 100 for best quality.758 Default: (90, 100).759 760 Returns:761 (Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],762 float32.763 """764 quality = np.random.uniform(quality_range[0], quality_range[1])765 return add_jpg_compression(img, int(quality))766