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

Feature/select noise models

noise_model.py

@@ -1,4 +1,6 @@
+import argparse
 import numpy as np
+import cv2
 
 
 def get_noise_model(noise_type="gaussian,0,50"):
@@ -18,3 +20,33 @@ def gaussian_noise(img):
         return gaussian_noise
     elif tokens[0] == "clean":
         return lambda img: img
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="test noise model",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--image_size", type=int, default=256,
+                        help="training patch size")
+    parser.add_argument("--noise_model", type=str, default="gaussian,0,50",
+                        help="noise model to be tested")
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+    image_size = args.image_size
+    noise_model = get_noise_model(args.noise_model)
+
+    while True:
+        image = np.ones((image_size, image_size, 3), dtype=np.uint8) * 128
+        cv2.imshow("noise image", noise_model(image))
+        key = cv2.waitKey(-1)
+
+        # "q": quit
+        if key == 113:
+            return 0
+
+
+if __name__ == '__main__':
+    main()

train.py

@@ -43,9 +43,9 @@ def get_args():
     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")
+                        help="noise model for source images")
     parser.add_argument("--target_noise_model", type=str, default="gaussian,25,25",
-                        help="checkpoint dir")
+                        help="noise model for target images")
     args = parser.parse_args()
 
     return args