predict_noisy_image_sdv_lab.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. from sklearn.externals import joblib
  2. import numpy as np
  3. from ipfml import image_processing
  4. from PIL import Image
  5. import sys, os, getopt
  6. min_max_file_path = './fichiersSVD_light/min_max_values'
  7. def main():
  8. if len(sys.argv) <= 1:
  9. print('Run with default parameters...')
  10. print('python predict_noisy_image.py --image path/to/xxxx --interval "0,20" --model path/to/xxxx.joblib --mode ["svdn", "svdne"]')
  11. sys.exit(2)
  12. try:
  13. opts, args = getopt.getopt(sys.argv[1:], "hi:t:m:o", ["help=", "image=", "interval=", "model=", "mode="])
  14. except getopt.GetoptError:
  15. # print help information and exit:
  16. print('python predict_noisy_image.py --image path/to/xxxx --interval "xx,xx" --model path/to/xxxx.joblib --mode ["svdn", "svdne"]')
  17. sys.exit(2)
  18. for o, a in opts:
  19. if o == "-h":
  20. print('python predict_noisy_image.py --image path/to/xxxx --interval "xx,xx" --model path/to/xxxx.joblib --mode ["svdn", "svdne"]')
  21. sys.exit()
  22. elif o in ("-i", "--image"):
  23. p_img_file = a
  24. elif o in ("-t", "--interval"):
  25. p_interval = list(map(int, a.split(',')))
  26. elif o in ("-m", "--model"):
  27. p_model_file = a
  28. elif o in ("-o", "--mode"):
  29. p_mode = a
  30. if p_mode != 'svdn' and p_mode != 'svdne' and p_mode != 'svd':
  31. assert False, "Mode not recognized"
  32. else:
  33. assert False, "unhandled option"
  34. # load of model file
  35. model = joblib.load(p_model_file)
  36. # load image
  37. img = Image.open(p_img_file)
  38. LAB_L = image_processing.get_LAB_L_SVD_s(img)
  39. # check mode to normalize data
  40. if p_mode == 'svdne':
  41. # need to read min_max_file
  42. with open(min_max_file_path, 'r') as f:
  43. min = float(f.readline().replace('\n', ''))
  44. max = float(f.readline().replace('\n', ''))
  45. l_values = image_processing.normalize_arr_with_range(LAB_L, min, max)
  46. elif p_mode == 'svdn':
  47. l_values = image_processing.normalize_arr(LAB_L)
  48. else:
  49. l_values = LAB_L
  50. # get interval values
  51. begin, end = p_interval
  52. test_data = l_values[begin:end]
  53. # get prediction of model
  54. prediction = model.predict([test_data])[0]
  55. print(prediction)
  56. if __name__== "__main__":
  57. main()