noise_computation.py 2.9 KB

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