mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Don't apply jpeg corruption & noise corruption together

Browse Source 
This causes some severe noise.
This commit is contained in:
James Betker 2020-10-20 12:56:35 -06:00
parent 111450f4e7
commit aba83e7497
2 changed files with 32 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
codes/data/image_corruptor.py
View File

@ -29,16 +29,17 @@ class ImageCorruptor:
rand_int_a = random.randint(1, 999999)
corrupted_imgs = []
applied_augs = augmentations + self.fixed_corruptions
for img in imgs:
for aug in augmentations:
img = self.apply_corruption(img, aug, rand_int_a)
img = self.apply_corruption(img, aug, rand_int_a, applied_augs)
for aug in self.fixed_corruptions:
img = self.apply_corruption(img, aug, rand_int_f)
img = self.apply_corruption(img, aug, rand_int_f, applied_augs)
corrupted_imgs.append(img)
return corrupted_imgs
def apply_corruption(self, img, aug, rand_int):
def apply_corruption(self, img, aug, rand_int, applied_augmentations):
if 'color_quantization' in aug:
# Color quantization
quant_div = 2 ** ((rand_int % 3) + 2)
@ -89,28 +90,31 @@ class ImageCorruptor:
noise_intensity = (rand_int % 4 + 2) / 255.0 # Between 1-4
img += np.random.randn(*img.shape) * noise_intensity
elif 'jpeg' in aug:
if aug == 'jpeg':
lo=10
range=20
elif aug == 'jpeg-medium':
lo=23
range=25
elif aug == 'jpeg-broad':
lo=15
range=60
# JPEG compression
qf = (rand_int % range + lo)
# cv2's jpeg compression is "odd". It introduces artifacts. Use PIL instead.
img = (img * 255).astype(np.uint8)
img = Image.fromarray(img)
buffer = BytesIO()
img.save(buffer, "JPEG", quality=qf, optimice=True)
buffer.seek(0)
jpeg_img_bytes = np.asarray(bytearray(buffer.read()), dtype="uint8")
img = read_img("buffer", jpeg_img_bytes, rgb=True)
if 'noise' not in applied_augmentations and 'noise-5' not in applied_augmentations:
if aug == 'jpeg':
lo=10
range=20
elif aug == 'jpeg-medium':
lo=23
range=25
elif aug == 'jpeg-broad':
lo=15
range=60
# JPEG compression
qf = (rand_int % range + lo)
# cv2's jpeg compression is "odd". It introduces artifacts. Use PIL instead.
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)
elif 'saturation' in aug:
# Lightening / saturation
saturation = float(rand_int % 10) * .03
img = np.clip(img + saturation, a_max=1, a_min=0)
elif 'none' not in aug:
raise NotImplementedError("Augmentation doesn't exist")
return img

11
codes/data/single_image_dataset.py
View File

@ -1,3 +1,4 @@
import random
from bisect import bisect_left
import numpy as np
import torch
@ -43,14 +44,14 @@ class SingleImageDataset(BaseUnsupervisedImageDataset):
if __name__ == '__main__':
opt = {
'name': 'amalgam',
'paths': ['F:\\4k6k\\datasets\\images\\flickr\\testbed'],
'paths': ['F:\\4k6k\\datasets\\images\\mi1_256'],
'weights': [1],
'target_size': 128,
'force_multiple': 32,
'scale': 2,
'eval': False,
'fixed_corruptions': ['jpeg'],
'random_corruptions': ['color_quantization', 'gaussian_blur', 'motion_blur', 'smooth_blur', 'noise', 'saturation'],
'fixed_corruptions': ['jpeg-broad'],
'random_corruptions': ['noise-5', 'none'],
'num_corrupts_per_image': 1
}
@ -58,9 +59,9 @@ if __name__ == '__main__':
import os
os.makedirs("debug", exist_ok=True)
for i in range(0, len(ds)):
o = ds[i]
o = ds[random.randint(0, len(ds))]
for k, v in o.items():
if 'path' not in k and 'center' not in k:
if 'LQ' in k and 'path' not in k and 'center' not in k:
#if 'full' in k:
#masked = v[:3, :, :] * v[3]
#torchvision.utils.save_image(masked.unsqueeze(0), "debug/%i_%s_masked.png" % (i, k))