noise_computation.py 3.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  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_param = None
  27. p_all = False
  28. if len(sys.argv) < 1:
  29. print('python noise_computation.py --noise xxxx --image path/to/image.png --n 100 --identical 0 --output image_name --all 1 --p 0.1')
  30. sys.exit(2)
  31. try:
  32. opts, args = getopt.getopt(sys.argv[1:], "h:n:i:n:i:o:a:p", ["help=", "noise=", "image=", "n=", "identical=", "output=", "all=", "p="])
  33. except getopt.GetoptError:
  34. # print help information and exit:
  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(2)
  37. for o, a in opts:
  38. if o == "-h":
  39. print('python noise_computation.py --noise xxxx --image path/to/image.png --n 100 --identical 0 --output image_name --all 1 --p 0.1')
  40. sys.exit()
  41. elif o in ("-n", "--noise"):
  42. p_noise = a
  43. if not p_noise in noise_list:
  44. assert False, "Unknow noise parameter %s, %s " % (p_noise, noise_list)
  45. elif o in ("-i", "--image"):
  46. p_image_path = a
  47. elif o in ("-n", "--n"):
  48. p_n = int(a)
  49. elif o in ("-i", "--identical"):
  50. p_identical = int(a)
  51. elif o in ("-o", "--output"):
  52. p_output = a
  53. elif o in ("-a", "--all"):
  54. p_all = int(a)
  55. elif o in ("-p", "--p"):
  56. p_param = float(a)
  57. else:
  58. assert False, "unhandled option"
  59. img = Image.open(p_image_path)
  60. if p_all:
  61. split_output = p_output.split('.')
  62. for i in range(1, p_n):
  63. p_filename = split_output[0] + "_" + str(i) + "." + filename_ext
  64. generate_noisy_image(img, i, p_noise, p_identical, p_filename, p_param)
  65. else:
  66. generate_noisy_image(img, p_n, p_noise, p_identical, p_output, p_param)
  67. if __name__== "__main__":
  68. main()