From aba83e7497d655860417a8c4f6aa9dcea584cc8e Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Tue, 20 Oct 2020 12:56:35 -0600
Subject: [PATCH] Don't apply jpeg corruption & noise corruption together

This causes some severe noise.
---
 codes/data/image_corruptor.py      | 48 ++++++++++++++++--------------
 codes/data/single_image_dataset.py | 11 +++----
 2 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/codes/data/image_corruptor.py b/codes/data/image_corruptor.py
index 317b4a7d..072ca7b0 100644
--- a/codes/data/image_corruptor.py
+++ b/codes/data/image_corruptor.py
@@ -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
diff --git a/codes/data/single_image_dataset.py b/codes/data/single_image_dataset.py
index 3bf5b0cb..855b00fa 100644
--- a/codes/data/single_image_dataset.py
+++ b/codes/data/single_image_dataset.py
@@ -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))