noise_computation.py 3.2 KB

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