noise_computation.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. from ipfml.filters import noise as nf
  2. import sys, os, getopt
  3. from PIL import Image
  4. from modules.utils import config as cfg
  5. from modules import noise
  6. noise_list = cfg.noise_labels
  7. generated_folder = cfg.generated_folder
  8. filename_ext = cfg.filename_ext
  9. def generate_noisy_image(p_image, p_n, p_noise, p_identical, p_output, p_param):
  10. noisy_image = noise.get_noise_result(p_image, p_n, _noise_choice=p_noise, _identical=p_identical, _p=p_param)
  11. noisy_image = Image.fromarray(noisy_image)
  12. output_path = os.path.join(generated_folder, p_noise)
  13. if not os.path.exists(output_path):
  14. os.makedirs(output_path)
  15. output_image_path = os.path.join(output_path, p_output)
  16. if not filename_ext in output_image_path:
  17. output_image_path = output_image_path + filename_ext
  18. noisy_image.save(output_image_path)
  19. print("Image saved at... '%s'" % output_image_path)
  20. def main():
  21. # by default..
  22. p_param = None
  23. p_all = False
  24. if len(sys.argv) < 1:
  25. print('python noise_computation.py --noise xxxx --image path/to/image.png --n 100 --identical 0 --output image_name --all 1 --p 0.1')
  26. sys.exit(2)
  27. try:
  28. opts, args = getopt.getopt(sys.argv[1:], "h:n:i:n:i:o:a:p", ["help=", "noise=", "image=", "n=", "identical=", "output=", "all=", "p="])
  29. except getopt.GetoptError:
  30. # print help information and exit:
  31. print('python noise_computation.py --noise xxxx --image path/to/image.png --n 100 --identical 0 --output image_name --all 1 --p 0.1')
  32. sys.exit(2)
  33. for o, a in opts:
  34. if o == "-h":
  35. print('python noise_computation.py --noise xxxx --image path/to/image.png --n 100 --identical 0 --output image_name --all 1 --p 0.1')
  36. sys.exit()
  37. elif o in ("-n", "--noise"):
  38. p_noise = a
  39. if not p_noise in noise_list:
  40. assert False, "Unknow noise parameter %s, %s " % (p_noise, noise_list)
  41. elif o in ("-i", "--image"):
  42. p_image_path = a
  43. elif o in ("-n", "--n"):
  44. p_n = int(a)
  45. elif o in ("-i", "--identical"):
  46. p_identical = int(a)
  47. elif o in ("-o", "--output"):
  48. p_output = a
  49. elif o in ("-a", "--all"):
  50. p_all = int(a)
  51. elif o in ("-p", "--p"):
  52. p_param = float(a)
  53. else:
  54. assert False, "unhandled option"
  55. img = Image.open(p_image_path)
  56. if p_all:
  57. split_output = p_output.split('.')
  58. for i in range(1, p_n):
  59. p_filename = split_output[0] + "_" + str(i) + "." + filename_ext
  60. generate_noisy_image(img, i, p_noise, p_identical, p_filename, p_param)
  61. else:
  62. generate_noisy_image(img, p_n, p_noise, p_identical, p_output, p_param)
  63. if __name__== "__main__":
  64. main()