Skip to content

Commit

Permalink
[Features] Support stochastic degradations for RealBasicVSR (open-mml…
Browse files Browse the repository at this point in the history 
…ab#647)

* support stochastic degradations

* fix bug in RandomBlur

* Fix bug in _apply_random_poisson_noise
  • Loading branch information
@ckkelvinchan
ckkelvinchan authored Dec 19, 2021
1 parent 64e0636 commit c00dc0d
Show file tree
Hide file tree
Showing 2 changed files with 122 additions and 47 deletions.
151 changes: 107 additions & 44 deletions mmedit/datasets/pipelines/random_degradations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,50 +31,76 @@ def __init__(self, params, keys):
self.keys = keys
self.params = params

def get_kernel(self):
def get_kernel(self, num_kernels):
kernel = np.random.choice(
self.params['kernel_list'], p=self.params['kernel_prob'])
kernel_size = random.choice(self.params['kernel_size'])

sigma_x = self.params.get('sigma_x', [0, 0])
sigma_x = np.random.uniform(sigma_x[0], sigma_x[1])
sigma_x_range = self.params.get('sigma_x', [0, 0])
sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
sigma_x_step = self.params.get('sigma_x_step', 0)

sigma_y = self.params.get('sigma_y', [0, 0])
sigma_y = np.random.uniform(sigma_y[0], sigma_y[1])
sigma_y_range = self.params.get('sigma_y', [0, 0])
sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
sigma_y_step = self.params.get('sigma_y_step', 0)

rotate_angle = self.params.get('rotate_angle', [-np.pi, np.pi])
rotate_angle = np.random.uniform(rotate_angle[0], rotate_angle[1])
rotate_angle_range = self.params.get('rotate_angle', [-np.pi, np.pi])
rotate_angle = np.random.uniform(rotate_angle_range[0],
rotate_angle_range[1])
rotate_angle_step = self.params.get('rotate_angle_step', 0)

beta_gau = self.params.get('beta_gaussian', [0.5, 4])
beta_gau = np.random.uniform(beta_gau[0], beta_gau[1])
beta_gau_range = self.params.get('beta_gaussian', [0.5, 4])
beta_gau = np.random.uniform(beta_gau_range[0], beta_gau_range[1])
beta_gau_step = self.params.get('beta_gaussian_step', 0)

beta_pla = self.params.get('beta_plateau', [1, 2])
beta_pla = np.random.uniform(beta_pla[0], beta_pla[1])
beta_pla_range = self.params.get('beta_plateau', [1, 2])
beta_pla = np.random.uniform(beta_pla_range[0], beta_pla_range[1])
beta_pla_step = self.params.get('beta_plateau_step', 0)

omega = self.params.get('omega', None)
if omega is not None:
omega = np.random.uniform(omega[0], omega[1])
else: # follow Real-ESRGAN
omega_range = self.params.get('omega', None)
omega_step = self.params.get('omega_step', 0)
if omega_range is None: # follow Real-ESRGAN settings if not specified
if kernel_size < 13:
omega = np.random.uniform(np.pi / 3, np.pi)
omega_range = [np.pi / 3., np.pi]
else:
omega = np.random.uniform(np.pi / 5, np.pi)
omega_range = [np.pi / 5., np.pi]
omega = np.random.uniform(omega_range[0], omega_range[1])

# determine blurring kernel
kernel = blur_kernels.random_mixed_kernels(
[kernel],
[1],
kernel_size,
[sigma_x, sigma_x],
[sigma_y, sigma_y],
[rotate_angle, rotate_angle],
[beta_gau, beta_gau],
[beta_pla, beta_pla],
[omega, omega],
None,
)

return kernel
kernels = []
for _ in range(0, num_kernels):
kernel = blur_kernels.random_mixed_kernels(
[kernel],
[1],
kernel_size,
[sigma_x, sigma_x],
[sigma_y, sigma_y],
[rotate_angle, rotate_angle],
[beta_gau, beta_gau],
[beta_pla, beta_pla],
[omega, omega],
None,
)
kernels.append(kernel)

# update kernel parameters
sigma_x += np.random.uniform(-sigma_x_step, sigma_x_step)
sigma_y += np.random.uniform(-sigma_y_step, sigma_y_step)
rotate_angle += np.random.uniform(-rotate_angle_step,
rotate_angle_step)
beta_gau += np.random.uniform(-beta_gau_step, beta_gau_step)
beta_pla += np.random.uniform(-beta_pla_step, beta_pla_step)
omega += np.random.uniform(-omega_step, omega_step)

sigma_x = np.clip(sigma_x, sigma_x_range[0], sigma_x_range[1])
sigma_y = np.clip(sigma_y, sigma_y_range[0], sigma_y_range[1])
rotate_angle = np.clip(rotate_angle, rotate_angle_range[0],
rotate_angle_range[1])
beta_gau = np.clip(beta_gau, beta_gau_range[0], beta_gau_range[1])
beta_pla = np.clip(beta_pla, beta_pla_range[0], beta_pla_range[1])
omega = np.clip(omega, omega_range[0], omega_range[1])

return kernels

def _apply_random_blur(self, imgs):
is_single_image = False
Expand All @@ -83,8 +109,11 @@ def _apply_random_blur(self, imgs):
imgs = [imgs]

# get kernel and blur the input
kernel = self.get_kernel()
imgs = [cv2.filter2D(im, -1, kernel) for im in imgs]
kernels = self.get_kernel(num_kernels=len(imgs))
imgs = [
cv2.filter2D(img, -1, kernel)
for img, kernel in zip(imgs, kernels)
]

if is_single_image:
imgs = imgs[0]
Expand Down Expand Up @@ -144,6 +173,8 @@ def _random_resize(self, imgs):
'implemented')
resize_opt = self.resize_dict[resize_opt]

resize_step = self.params.get('resize_step', 0)

# determine the target size, if not provided
target_size = self.params.get('target_size', None)
if target_size is None:
Expand All @@ -157,16 +188,32 @@ def _random_resize(self, imgs):
else:
scale_factor = 1
target_size = (int(h * scale_factor), int(w * scale_factor))
else:
resize_step = 0

# resize the input
imgs = [
cv2.resize(img, target_size[::-1], interpolation=resize_opt)
for img in imgs
]
if resize_step == 0: # same target_size for all input images
outputs = [
cv2.resize(img, target_size[::-1], interpolation=resize_opt)
for img in imgs
]
else: # different target_size for each input image
outputs = []
for img in imgs:
img = cv2.resize(
img, target_size[::-1], interpolation=resize_opt)
outputs.append(img)

# update scale
scale_factor += np.random.uniform(-resize_step, resize_step)
scale_factor = np.clip(scale_factor, resize_scale[0],
resize_scale[1])
target_size = (int(h * scale_factor), int(w * scale_factor))

if is_single_image:
imgs = imgs[0]
outputs = outputs[0]

return imgs
return outputs

def __call__(self, results):
if np.random.uniform() > self.params.get('prob', 1):
Expand Down Expand Up @@ -205,6 +252,8 @@ def _apply_gaussian_noise(self, imgs):
sigma_range = self.params['gaussian_sigma']
sigma = np.random.uniform(sigma_range[0], sigma_range[1]) / 255.

sigma_step = self.params.get('gaussian_sigma_step', 0)

gray_noise_prob = self.params['gaussian_gray_noise_prob']
is_gray_noise = np.random.uniform() < gray_noise_prob

Expand All @@ -215,12 +264,19 @@ def _apply_gaussian_noise(self, imgs):
noise = noise[:, :, :1]
outputs.append(img + noise)

# update noise level
sigma += np.random.uniform(-sigma_step, sigma_step) / 255.
sigma = np.clip(sigma, sigma_range[0] / 255.,
sigma_range[1] / 255.)

return outputs

def _apply_poisson_noise(self, imgs):
scale_range = self.params['poisson_scale']
scale = np.random.uniform(scale_range[0], scale_range[1])

scale_step = self.params.get('poisson_scale_step', 0)

gray_noise_prob = self.params['poisson_gray_noise_prob']
is_gray_noise = np.random.uniform() < gray_noise_prob

Expand All @@ -236,6 +292,10 @@ def _apply_poisson_noise(self, imgs):

outputs.append(img + noise * scale)

# update noise level
scale += np.random.uniform(-scale_step, scale_step)
scale = np.clip(scale, scale_range[0], scale_range[1])

return outputs

def _apply_random_noise(self, imgs):
Expand Down Expand Up @@ -297,23 +357,26 @@ def _apply_random_compression(self, imgs):
is_single_image = True
imgs = [imgs]

# determine compression level
# determine initial compression level and the step size
quality = self.params['quality']
quality_step = self.params.get('quality_step', 0)
jpeg_param = round(np.random.uniform(quality[0], quality[1]))
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_param]

# apply jpeg compression
outputs = []
for img in imgs:
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_param]
_, img_encoded = cv2.imencode('.jpg', img * 255., encode_param)
outputs.append(np.float32(cv2.imdecode(img_encoded, 1)) / 255.)

imgs = outputs
# update compression level
jpeg_param += np.random.uniform(-quality_step, quality_step)
jpeg_param = round(np.clip(jpeg_param, quality[0], quality[1]))

if is_single_image:
imgs = imgs[0]
outputs = outputs[0]

return imgs
return outputs

def __call__(self, results):
if np.random.uniform() > self.params.get('prob', 1):
Expand Down
18 changes: 15 additions & 3 deletions tests/test_data/test_pipelines/test_random_degradations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,12 +138,24 @@ def test_random_resize():
resize_scale=[0.5, 1.5],
resize_opt=['bilinear', 'area', 'bicubic'],
resize_prob=[1 / 3., 1 / 3., 1 / 3.],
target_size=(16, 16)),
target_size=(16, 32)),
keys=['lq'])
results = model(results)
assert results['lq'].shape == (16, 16, 3)
assert results['lq'].shape == (16, 32, 3)

# skip degrdation
# step_size > 0
results['lq'] = np.ones((8, 8, 3)).astype(np.float32)
model = RandomResize(
params=dict(
resize_mode_prob=[0, 0, 1],
resize_scale=[0.5, 1.5],
resize_opt=['bilinear', 'area', 'bicubic'],
resize_prob=[1 / 3., 1 / 3., 1 / 3.],
resize_step=0.05),
keys=['lq'])
results = model(results)

# skip degradation
model = RandomResize(
params=dict(
resize_mode_prob=[1, 0, 0],
Expand Down

0 comments on commit c00dc0d

Please sign in to comment.