2020-09-18 15:49:13 +00:00
|
|
|
import random
|
2021-06-07 15:13:54 +00:00
|
|
|
from math import cos, pi
|
|
|
|
|
2020-09-25 22:37:54 +00:00
|
|
|
import cv2
|
|
|
|
import numpy as np
|
|
|
|
from data.util import read_img
|
|
|
|
from PIL import Image
|
|
|
|
from io import BytesIO
|
2020-09-18 15:49:13 +00:00
|
|
|
|
2021-06-07 15:13:54 +00:00
|
|
|
|
|
|
|
# Feeds a random uniform through a cosine distribution to slightly bias corruptions towards "uncorrupted".
|
|
|
|
# Return is on [0,1] with a bias towards 0.
|
|
|
|
def get_rand():
|
|
|
|
r = random.random()
|
|
|
|
return 1 - cos(r * pi / 2)
|
|
|
|
|
|
|
|
# Get a rough visualization of the above distribution. (Y-axis is meaningless, just spreads data)
|
|
|
|
'''
|
|
|
|
if __name__ == '__main__':
|
|
|
|
import numpy as np
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
data = np.asarray([get_rand() for _ in range(5000)])
|
|
|
|
plt.plot(data, np.random.uniform(size=(5000,)), 'x')
|
|
|
|
plt.show()
|
|
|
|
'''
|
|
|
|
|
2020-09-18 15:49:13 +00:00
|
|
|
# Performs image corruption on a list of images from a configurable set of corruption
|
|
|
|
# options.
|
|
|
|
class ImageCorruptor:
|
|
|
|
def __init__(self, opt):
|
2020-12-26 20:49:27 +00:00
|
|
|
self.blur_scale = opt['corruption_blur_scale'] if 'corruption_blur_scale' in opt.keys() else 1
|
2020-12-06 03:30:36 +00:00
|
|
|
self.fixed_corruptions = opt['fixed_corruptions'] if 'fixed_corruptions' in opt.keys() else []
|
2020-11-25 20:59:06 +00:00
|
|
|
self.num_corrupts = opt['num_corrupts_per_image'] if 'num_corrupts_per_image' in opt.keys() else 0
|
2020-09-26 04:19:38 +00:00
|
|
|
if self.num_corrupts == 0:
|
|
|
|
return
|
2020-11-25 20:59:06 +00:00
|
|
|
self.random_corruptions = opt['random_corruptions'] if 'random_corruptions' in opt.keys() else []
|
2020-09-18 15:49:13 +00:00
|
|
|
|
2021-06-07 15:13:54 +00:00
|
|
|
def corrupt_images(self, imgs, return_entropy=False):
|
2020-09-28 20:26:15 +00:00
|
|
|
if self.num_corrupts == 0 and not self.fixed_corruptions:
|
2021-06-11 21:31:10 +00:00
|
|
|
if return_entropy:
|
|
|
|
return imgs, []
|
|
|
|
else:
|
|
|
|
return imgs
|
2020-09-26 04:19:38 +00:00
|
|
|
|
2020-09-28 20:26:15 +00:00
|
|
|
if self.num_corrupts == 0:
|
|
|
|
augmentations = []
|
|
|
|
else:
|
|
|
|
augmentations = random.choices(self.random_corruptions, k=self.num_corrupts)
|
2020-09-18 15:49:13 +00:00
|
|
|
|
2021-06-07 15:13:54 +00:00
|
|
|
# Sources of entropy
|
2020-09-18 15:49:13 +00:00
|
|
|
corrupted_imgs = []
|
2021-06-07 15:13:54 +00:00
|
|
|
entropy = []
|
2020-10-20 18:56:35 +00:00
|
|
|
applied_augs = augmentations + self.fixed_corruptions
|
2020-09-18 15:49:13 +00:00
|
|
|
for img in imgs:
|
|
|
|
for aug in augmentations:
|
2021-06-07 15:13:54 +00:00
|
|
|
r = get_rand()
|
|
|
|
img = self.apply_corruption(img, aug, r, applied_augs)
|
2020-09-26 04:45:57 +00:00
|
|
|
for aug in self.fixed_corruptions:
|
2021-06-07 15:13:54 +00:00
|
|
|
r = get_rand()
|
|
|
|
img = self.apply_corruption(img, aug, r, applied_augs)
|
|
|
|
entropy.append(r)
|
2020-09-25 22:37:54 +00:00
|
|
|
corrupted_imgs.append(img)
|
|
|
|
|
2021-06-07 15:13:54 +00:00
|
|
|
if return_entropy:
|
|
|
|
return corrupted_imgs, entropy
|
|
|
|
else:
|
|
|
|
return corrupted_imgs
|
2020-09-18 15:49:13 +00:00
|
|
|
|
2021-06-07 15:13:54 +00:00
|
|
|
def apply_corruption(self, img, aug, rand_val, applied_augmentations):
|
2020-09-25 22:37:54 +00:00
|
|
|
if 'color_quantization' in aug:
|
|
|
|
# Color quantization
|
2021-06-07 15:13:54 +00:00
|
|
|
quant_div = 2 ** (int(rand_val * 10 / 3) + 2)
|
2020-09-25 22:37:54 +00:00
|
|
|
img = img * 255
|
|
|
|
img = (img // quant_div) * quant_div
|
|
|
|
img = img / 255
|
|
|
|
elif 'gaussian_blur' in aug:
|
2021-06-07 15:13:54 +00:00
|
|
|
img = cv2.GaussianBlur(img, (0,0), rand_val*1.5)
|
2020-09-25 22:37:54 +00:00
|
|
|
elif 'motion_blur' in aug:
|
|
|
|
# Motion blur
|
2021-06-07 15:13:54 +00:00
|
|
|
intensity = self.blur_scale * rand_val * 3 + 1
|
|
|
|
angle = random.randint(0,360)
|
2020-09-25 22:37:54 +00:00
|
|
|
k = np.zeros((intensity, intensity), dtype=np.float32)
|
|
|
|
k[(intensity - 1) // 2, :] = np.ones(intensity, dtype=np.float32)
|
|
|
|
k = cv2.warpAffine(k, cv2.getRotationMatrix2D((intensity / 2 - 0.5, intensity / 2 - 0.5), angle, 1.0),
|
|
|
|
(intensity, intensity))
|
|
|
|
k = k * (1.0 / np.sum(k))
|
|
|
|
img = cv2.filter2D(img, -1, k)
|
|
|
|
elif 'block_noise' in aug:
|
|
|
|
# Large distortion blocks in part of an img, such as is used to mask out a face.
|
|
|
|
pass
|
|
|
|
elif 'lq_resampling' in aug:
|
2020-10-31 17:08:55 +00:00
|
|
|
# Random mode interpolation HR->LR->HR
|
|
|
|
scale = 2
|
|
|
|
if 'lq_resampling4x' == aug:
|
|
|
|
scale = 4
|
2020-12-26 20:49:27 +00:00
|
|
|
interpolation_modes = [cv2.INTER_NEAREST, cv2.INTER_CUBIC, cv2.INTER_LINEAR, cv2.INTER_LANCZOS4]
|
2021-06-07 15:13:54 +00:00
|
|
|
mode = random.randint(0,4) % len(interpolation_modes)
|
2020-10-31 17:08:55 +00:00
|
|
|
# Downsample first, then upsample using the random mode.
|
|
|
|
img = cv2.resize(img, dsize=(img.shape[1]//scale, img.shape[0]//scale), interpolation=cv2.INTER_NEAREST)
|
|
|
|
img = cv2.resize(img, dsize=(img.shape[1]*scale, img.shape[0]*scale), interpolation=mode)
|
2020-09-25 22:37:54 +00:00
|
|
|
elif 'color_shift' in aug:
|
|
|
|
# Color shift
|
|
|
|
pass
|
|
|
|
elif 'interlacing' in aug:
|
|
|
|
# Interlacing distortion
|
|
|
|
pass
|
|
|
|
elif 'chromatic_aberration' in aug:
|
|
|
|
# Chromatic aberration
|
|
|
|
pass
|
|
|
|
elif 'noise' in aug:
|
2020-10-15 16:12:25 +00:00
|
|
|
# Random noise
|
|
|
|
if 'noise-5' == aug:
|
|
|
|
noise_intensity = 5 / 255.0
|
|
|
|
else:
|
2021-06-07 15:13:54 +00:00
|
|
|
noise_intensity = (rand_val*4 + 2) / 255.0
|
2020-09-27 17:12:24 +00:00
|
|
|
img += np.random.randn(*img.shape) * noise_intensity
|
2020-09-25 22:37:54 +00:00
|
|
|
elif 'jpeg' in aug:
|
2020-10-20 18:56:35 +00:00
|
|
|
if 'noise' not in applied_augmentations and 'noise-5' not in applied_augmentations:
|
|
|
|
if aug == 'jpeg':
|
|
|
|
lo=10
|
|
|
|
range=20
|
2020-11-25 20:59:06 +00:00
|
|
|
elif aug == 'jpeg-low':
|
|
|
|
lo=15
|
|
|
|
range=10
|
2020-10-20 18:56:35 +00:00
|
|
|
elif aug == 'jpeg-medium':
|
|
|
|
lo=23
|
|
|
|
range=25
|
|
|
|
elif aug == 'jpeg-broad':
|
|
|
|
lo=15
|
|
|
|
range=60
|
2020-10-28 22:40:47 +00:00
|
|
|
elif aug == 'jpeg-normal':
|
2020-11-01 02:50:24 +00:00
|
|
|
lo=47
|
2020-10-28 23:37:39 +00:00
|
|
|
range=35
|
|
|
|
else:
|
|
|
|
raise NotImplementedError("specified jpeg corruption doesn't exist")
|
2020-10-20 18:56:35 +00:00
|
|
|
# JPEG compression
|
2021-06-07 15:13:54 +00:00
|
|
|
qf = (int((1-rand_val)*range) + lo)
|
2020-12-31 17:13:24 +00:00
|
|
|
# Use PIL to perform a mock compression to a data buffer, then swap back to cv2.
|
2020-10-20 18:56:35 +00:00
|
|
|
img = (img * 255).astype(np.uint8)
|
|
|
|
img = Image.fromarray(img)
|
|
|
|
buffer = BytesIO()
|
|
|
|
img.save(buffer, "JPEG", quality=qf, optimize=True)
|
|
|
|
buffer.seek(0)
|
|
|
|
jpeg_img_bytes = np.asarray(bytearray(buffer.read()), dtype="uint8")
|
|
|
|
img = read_img("buffer", jpeg_img_bytes, rgb=True)
|
2020-09-25 22:37:54 +00:00
|
|
|
elif 'saturation' in aug:
|
|
|
|
# Lightening / saturation
|
2021-06-07 15:13:54 +00:00
|
|
|
saturation = rand_val * .3
|
2020-09-25 22:37:54 +00:00
|
|
|
img = np.clip(img + saturation, a_max=1, a_min=0)
|
2020-10-20 18:56:35 +00:00
|
|
|
elif 'none' not in aug:
|
|
|
|
raise NotImplementedError("Augmentation doesn't exist")
|
2020-09-18 15:49:13 +00:00
|
|
|
|
2020-09-25 22:37:54 +00:00
|
|
|
return img
|