Merge pull request yu4u#1 from yu4u/feature/select_noise_models

Feature/select noise models

generator.py

@@ -5,24 +5,14 @@
 from keras.utils import Sequence
 
 
-def add_noise(img, min_stddev=0, max_stddev=50):
-    noise_img = img.astype(np.float)
-    stddev = np.random.uniform(min_stddev, max_stddev)
-    noise = np.random.randn(*img.shape) * stddev
-    noise_img += noise
-    noise_img = np.clip(noise_img, 0, 255).astype(np.uint8)
-
-    return noise_img
-
-
 class NoisyImageGenerator(Sequence):
-    def __init__(self, image_dir, batch_size=32, image_size=64, min_stddev=0, max_stddev=50):
+    def __init__(self, image_dir, source_noise_model, target_noise_model, batch_size=32, image_size=64):
         self.image_paths = list(Path(image_dir).glob("*.jpg"))
+        self.source_noise_model = source_noise_model
+        self.target_noise_model = target_noise_model
         self.image_num = len(self.image_paths)
         self.batch_size = batch_size
         self.image_size = image_size
-        self.min_stddev = min_stddev
-        self.max_stddev = max_stddev
 
     def __len__(self):
         return self.image_num // self.batch_size
@@ -44,9 +34,8 @@ def __getitem__(self, idx):
                 i = np.random.randint(h - image_size + 1)
                 j = np.random.randint(w - image_size + 1)
                 clean_patch = image[i:i + image_size, j:j + image_size]
-
-                y[sample_id] = add_noise(clean_patch, min_stddev=self.min_stddev, max_stddev=self.max_stddev)
-                x[sample_id] = add_noise(clean_patch, min_stddev=self.min_stddev, max_stddev=self.max_stddev)
+                x[sample_id] = self.source_noise_model(clean_patch)
+                y[sample_id] = self.target_noise_model(clean_patch)
 
                 sample_id += 1
 
@@ -55,14 +44,14 @@ def __getitem__(self, idx):
 
 
 class ValGenerator(Sequence):
-    def __init__(self, image_dir, stddev=25):
+    def __init__(self, image_dir, source_noise_model):
         image_paths = list(Path(image_dir).glob("*.*"))
         self.image_num = len(image_paths)
         self.data = []
 
         for image_path in image_paths:
             y = cv2.imread(str(image_path))
-            x = add_noise(y, min_stddev=stddev, max_stddev=stddev)
+            x = source_noise_model(y)
             self.data.append([np.expand_dims(x, axis=0), np.expand_dims(y, axis=0)])
 
     def __len__(self):

noise_model.py

@@ -0,0 +1,20 @@
+import numpy as np
+
+
+def get_noise_model(noise_type="gaussian,0,50"):
+    tokens = noise_type.split(sep=",")
+
+    if tokens[0] == "gaussian":
+        min_stddev = int(tokens[1])
+        max_stddev = int(tokens[2])
+
+        def gaussian_noise(img):
+            noise_img = img.astype(np.float)
+            stddev = np.random.uniform(min_stddev, max_stddev)
+            noise = np.random.randn(*img.shape) * stddev
+            noise_img += noise
+            noise_img = np.clip(noise_img, 0, 255).astype(np.uint8)
+            return noise_img
+        return gaussian_noise
+    elif tokens[0] == "clean":
+        return lambda img: img

train.py

@@ -5,6 +5,7 @@
 from keras.optimizers import Adam
 from model import get_srresnet_model, PSNR
 from generator import NoisyImageGenerator, ValGenerator
+from noise_model import get_noise_model
 
 
 class Schedule:
@@ -41,6 +42,10 @@ def get_args():
                         help="steps per epoch")
     parser.add_argument("--output_path", type=str, default="checkpoints",
                         help="checkpoint dir")
+    parser.add_argument("--source_noise_model", type=str, default="gaussian,0,50",
+                        help="checkpoint dir")
+    parser.add_argument("--target_noise_model", type=str, default="gaussian,25,25",
+                        help="checkpoint dir")
     args = parser.parse_args()
 
     return args
@@ -59,16 +64,19 @@ def main():
     model = get_srresnet_model()
     opt = Adam(lr=lr)
     model.compile(optimizer=opt, loss="mse", metrics=[PSNR])
-    generator = NoisyImageGenerator(image_dir, batch_size=batch_size, image_size=image_size)
-    val_generator = ValGenerator(test_dir)
+    source_noise_model = get_noise_model(args.source_noise_model)
+    target_noise_model = get_noise_model(args.target_noise_model)
+    generator = NoisyImageGenerator(image_dir, source_noise_model, target_noise_model, batch_size=batch_size,
+                                    image_size=image_size)
+    val_generator = ValGenerator(test_dir, source_noise_model)
     output_path.mkdir(parents=True, exist_ok=True)
     callbacks = [
         LearningRateScheduler(schedule=Schedule(nb_epochs, lr)),
         ModelCheckpoint(str(output_path) + "/weights.{epoch:03d}-{val_loss:.3f}-{val_PSNR:.5f}.hdf5",
                         monitor="val_PSNR",
                         verbose=1,
                         mode="max",
-                        save_best_only=True, )
+                        save_best_only=True)
     ]
 
     hist = model.fit_generator(generator=generator,