Skip to content

Enhance img_process_util.py: Add return types and improve USM sharpen… #732

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lunikdev wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Enhance img_process_util.py: Add return types and improve USM sharpen… #732

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
140 changes: 23 additions & 117 deletions basicsr/data/degradations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
# degradations.py

import cv2
import math
import numpy as np
import random
import torch
from scipy import special
from scipy.stats import multivariate_normal
from torchvision.transforms.functional import rgb_to_grayscale
from torchvision.transforms import functional as F

# -------------------------------------------------------------------- #
# --------------------------- blur kernels --------------------------- #
Expand Down Expand Up @@ -112,8 +114,6 @@ def bivariate_Gaussian(kernel_size, sig_x, sig_y, theta, grid=None, isotropic=Tr
def bivariate_generalized_Gaussian(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
"""Generate a bivariate generalized Gaussian kernel.

``Paper: Parameter Estimation For Multivariate Generalized Gaussian Distributions``

In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.

Args:
Expand Down Expand Up @@ -143,10 +143,6 @@ def bivariate_generalized_Gaussian(kernel_size, sig_x, sig_y, theta, beta, grid=
def bivariate_plateau(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
"""Generate a plateau-like anisotropic kernel.

1 / (1+x^(beta))

Reference: https://stats.stackexchange.com/questions/203629/is-there-a-plateau-shaped-distribution

In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.

Args:
Expand Down Expand Up @@ -179,21 +175,7 @@ def random_bivariate_Gaussian(kernel_size,
rotation_range,
noise_range=None,
isotropic=True):
"""Randomly generate bivariate isotropic or anisotropic Gaussian kernels.

In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.

Args:
kernel_size (int):
sigma_x_range (tuple): [0.6, 5]
sigma_y_range (tuple): [0.6, 5]
rotation range (tuple): [-math.pi, math.pi]
noise_range(tuple, optional): multiplicative kernel noise,
[0.75, 1.25]. Default: None

Returns:
kernel (ndarray):
"""
"""Randomly generate bivariate isotropic or anisotropic Gaussian kernels."""
assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
Expand Down Expand Up @@ -224,22 +206,7 @@ def random_bivariate_generalized_Gaussian(kernel_size,
beta_range,
noise_range=None,
isotropic=True):
"""Randomly generate bivariate generalized Gaussian kernels.

In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.

Args:
kernel_size (int):
sigma_x_range (tuple): [0.6, 5]
sigma_y_range (tuple): [0.6, 5]
rotation range (tuple): [-math.pi, math.pi]
beta_range (tuple): [0.5, 8]
noise_range(tuple, optional): multiplicative kernel noise,
[0.75, 1.25]. Default: None

Returns:
kernel (ndarray):
"""
"""Randomly generate bivariate generalized Gaussian kernels."""
assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
Expand All @@ -252,15 +219,13 @@ def random_bivariate_generalized_Gaussian(kernel_size,
sigma_y = sigma_x
rotation = 0

# assume beta_range[0] < 1 < beta_range[1]
if np.random.uniform() < 0.5:
beta = np.random.uniform(beta_range[0], 1)
else:
beta = np.random.uniform(1, beta_range[1])

kernel = bivariate_generalized_Gaussian(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)

# add multiplicative noise
if noise_range is not None:
assert noise_range[0] < noise_range[1], 'Wrong noise range.'
noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
Expand All @@ -276,22 +241,7 @@ def random_bivariate_plateau(kernel_size,
beta_range,
noise_range=None,
isotropic=True):
"""Randomly generate bivariate plateau kernels.

In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.

Args:
kernel_size (int):
sigma_x_range (tuple): [0.6, 5]
sigma_y_range (tuple): [0.6, 5]
rotation range (tuple): [-math.pi/2, math.pi/2]
beta_range (tuple): [1, 4]
noise_range(tuple, optional): multiplicative kernel noise,
[0.75, 1.25]. Default: None

Returns:
kernel (ndarray):
"""
"""Randomly generate bivariate plateau kernels."""
assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
Expand All @@ -304,14 +254,12 @@ def random_bivariate_plateau(kernel_size,
sigma_y = sigma_x
rotation = 0

# TODO: this may be not proper
if np.random.uniform() < 0.5:
beta = np.random.uniform(beta_range[0], 1)
else:
beta = np.random.uniform(1, beta_range[1])

kernel = bivariate_plateau(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)
# add multiplicative noise
if noise_range is not None:
assert noise_range[0] < noise_range[1], 'Wrong noise range.'
noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
Expand All @@ -330,25 +278,7 @@ def random_mixed_kernels(kernel_list,
betag_range=(0.5, 8),
betap_range=(0.5, 8),
noise_range=None):
"""Randomly generate mixed kernels.

Args:
kernel_list (tuple): a list name of kernel types,
support ['iso', 'aniso', 'skew', 'generalized', 'plateau_iso',
'plateau_aniso']
kernel_prob (tuple): corresponding kernel probability for each
kernel type
kernel_size (int):
sigma_x_range (tuple): [0.6, 5]
sigma_y_range (tuple): [0.6, 5]
rotation range (tuple): [-math.pi, math.pi]
beta_range (tuple): [0.5, 8]
noise_range(tuple, optional): multiplicative kernel noise,
[0.75, 1.25]. Default: None

Returns:
kernel (ndarray):
"""
"""Randomly generate mixed kernels."""
kernel_type = random.choices(kernel_list, kernel_prob)[0]
if kernel_type == 'iso':
kernel = random_bivariate_Gaussian(
Expand Down Expand Up @@ -383,9 +313,6 @@ def random_mixed_kernels(kernel_list,
return kernel


np.seterr(divide='ignore', invalid='ignore')


def circular_lowpass_kernel(cutoff, kernel_size, pad_to=0):
"""2D sinc filter

Expand Down Expand Up @@ -534,9 +461,9 @@ def random_add_gaussian_noise(img, sigma_range=(0, 1.0), gray_prob=0, clip=True,


def random_generate_gaussian_noise_pt(img, sigma_range=(0, 10), gray_prob=0):
sigma = torch.rand(
img.size(0), dtype=img.dtype, device=img.device) * (sigma_range[1] - sigma_range[0]) + sigma_range[0]
gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
device = img.device # Get the device of input tensor
sigma = torch.rand(img.size(0), dtype=img.dtype, device=device) * (sigma_range[1] - sigma_range[0]) + sigma_range[0]
gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=device)
gray_noise = (gray_noise < gray_prob).float()
return generate_gaussian_noise_pt(img, sigma, gray_noise)

Expand Down Expand Up @@ -607,46 +534,34 @@ def add_poisson_noise(img, scale=1.0, clip=True, rounds=False, gray_noise=False)


def generate_poisson_noise_pt(img, scale=1.0, gray_noise=0):
"""Generate a batch of poisson noise (PyTorch version)

Args:
img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.
scale (float | Tensor): Noise scale. Number or Tensor with shape (b).
Default: 1.0.
gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).
0 for False, 1 for True. Default: 0.

Returns:
(Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
float32.
"""
"""Generate a batch of poisson noise (PyTorch version)"""
b, _, h, w = img.size()
device = img.device # Get the device of input tensor

if isinstance(gray_noise, (float, int)):
cal_gray_noise = gray_noise > 0
else:
gray_noise = gray_noise.view(b, 1, 1, 1)
cal_gray_noise = torch.sum(gray_noise) > 0

if cal_gray_noise:
img_gray = rgb_to_grayscale(img, num_output_channels=1)
# round and clip image for counting vals correctly
img_gray = F.rgb_to_grayscale(img)
img_gray = torch.clamp((img_gray * 255.0).round(), 0, 255) / 255.
# use for-loop to get the unique values for each sample
vals_list = [len(torch.unique(img_gray[i, :, :, :])) for i in range(b)]
vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]
vals = img_gray.new_tensor(vals_list).view(b, 1, 1, 1)
vals = torch.tensor(vals_list, dtype=img.dtype, device=device).view(b, 1, 1, 1) # Create tensor on correct device
out = torch.poisson(img_gray * vals) / vals
noise_gray = out - img_gray
noise_gray = noise_gray.expand(b, 3, h, w)

# always calculate color noise
# round and clip image for counting vals correctly
# Color noise
img = torch.clamp((img * 255.0).round(), 0, 255) / 255.
# use for-loop to get the unique values for each sample
vals_list = [len(torch.unique(img[i, :, :, :])) for i in range(b)]
vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]
vals = img.new_tensor(vals_list).view(b, 1, 1, 1)
vals = torch.tensor(vals_list, dtype=img.dtype, device=device).view(b, 1, 1, 1) # Create tensor on correct device
out = torch.poisson(img * vals) / vals
noise = out - img

if cal_gray_noise:
noise = noise * (1 - gray_noise) + noise_gray * gray_noise
if not isinstance(scale, (float, int)):
Expand Down Expand Up @@ -682,15 +597,6 @@ def add_poisson_noise_pt(img, scale=1.0, clip=True, rounds=False, gray_noise=0):
# ----------------------- Random Poisson (Shot) Noise ----------------------- #


def random_generate_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0):
scale = np.random.uniform(scale_range[0], scale_range[1])
if np.random.uniform() < gray_prob:
gray_noise = True
else:
gray_noise = False
return generate_poisson_noise(img, scale, gray_noise)


def random_add_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
noise = random_generate_poisson_noise(img, scale_range, gray_prob)
out = img + noise
Expand All @@ -704,9 +610,9 @@ def random_add_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0, clip=True,


def random_generate_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0):
scale = torch.rand(
img.size(0), dtype=img.dtype, device=img.device) * (scale_range[1] - scale_range[0]) + scale_range[0]
gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
device = img.device # Get the device of input tensor
scale = torch.rand(img.size(0), dtype=img.dtype, device=device) * (scale_range[1] - scale_range[0]) + scale_range[0]
gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=device)
gray_noise = (gray_noise < gray_prob).float()
return generate_poisson_noise_pt(img, scale, gray_noise)

Expand Down Expand Up @@ -761,4 +667,4 @@ def random_add_jpg_compression(img, quality_range=(90, 100)):
float32.
"""
quality = np.random.uniform(quality_range[0], quality_range[1])
return add_jpg_compression(img, quality)
return add_jpg_compression(img, quality)
61 changes: 45 additions & 16 deletions basicsr/utils/img_process_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
# img_process_util.py

import cv2
import numpy as np
import torch
from torch.nn import functional as F
import torch.nn.functional as F


def filter2D(img, kernel):
Expand All @@ -10,6 +12,9 @@ def filter2D(img, kernel):
Args:
img (Tensor): (b, c, h, w)
kernel (Tensor): (b, k, k)

Returns:
Tensor: Filtered image with shape (b, c, h, w)
"""
k = kernel.size(-1)
b, c, h, w = img.size()
Expand All @@ -34,33 +39,44 @@ def filter2D(img, kernel):
def usm_sharp(img, weight=0.5, radius=50, threshold=10):
"""USM sharpening.

Input image: I; Blurry image: B.
1. sharp = I + weight * (I - B)
2. Mask = 1 if abs(I - B) > threshold, else: 0
3. Blur mask:
4. Out = Mask * sharp + (1 - Mask) * I


Args:
img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
img (Numpy array or Tensor): Input image, HWC, BGR; float32, [0, 1].
weight (float): Sharp weight. Default: 1.
radius (float): Kernel size of Gaussian blur. Default: 50.
threshold (int):
threshold (int): Threshold for the sharpening mask.

Returns:
Numpy array or Tensor: Sharpened image, same type as input
"""
if radius % 2 == 0:
radius += 1
blur = cv2.GaussianBlur(img, (radius, radius), 0)
residual = img - blur

if torch.is_tensor(img):
device = img.device
# Move to CPU for OpenCV operations
img_np = img.cpu().numpy()
else:
device = None
img_np = img

blur = cv2.GaussianBlur(img_np, (radius, radius), 0)
residual = img_np - blur
mask = np.abs(residual) * 255 > threshold
mask = mask.astype('float32')
soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)

sharp = img + weight * residual
sharp = img_np + weight * residual
sharp = np.clip(sharp, 0, 1)
return soft_mask * sharp + (1 - soft_mask) * img
output = soft_mask * sharp + (1 - soft_mask) * img_np

if device is not None:
# Convert back to tensor if input was tensor
output = torch.from_numpy(output).to(device)
return output


class USMSharp(torch.nn.Module):
"""PyTorch version of Unsharp Masking sharpening layer."""

def __init__(self, radius=50, sigma=0):
super(USMSharp, self).__init__()
Expand All @@ -72,12 +88,25 @@ def __init__(self, radius=50, sigma=0):
self.register_buffer('kernel', kernel)

def forward(self, img, weight=0.5, threshold=10):
"""Forward function.

Args:
img (Tensor): Input image tensor (b, c, h, w)
weight (float): Sharpening weight
threshold (float): Threshold for the sharpening mask

Returns:
Tensor: Sharpened image tensor (b, c, h, w)
"""
device = img.device
self.kernel = self.kernel.to(device) # Ensure kernel is on same device as input

blur = filter2D(img, self.kernel)
residual = img - blur

mask = torch.abs(residual) * 255 > threshold
mask = mask.float()
soft_mask = filter2D(mask, self.kernel)
sharp = img + weight * residual
sharp = torch.clip(sharp, 0, 1)
return soft_mask * sharp + (1 - soft_mask) * img
sharp = torch.clamp(sharp, 0, 1)
return soft_mask * sharp + (1 - soft_mask) * img