reconstruct_keras.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. import numpy as np
  2. import pandas as pd
  3. import json
  4. import os, sys, argparse
  5. from keras.models import model_from_json
  6. import modules.config as cfg
  7. from modules.features import compute_feature
  8. from joblib import dump, load
  9. from PIL import Image
  10. def reconstruct(_scene_name, _model_path, _n, _feature_choice):
  11. # construct the empty output image
  12. output_image = np.empty([cfg.number_of_rows, cfg.number_of_columns])
  13. # load the trained model
  14. with open(_model_path, 'r') as f:
  15. json_model = json.load(f)
  16. model = model_from_json(json_model)
  17. model.load_weights(_model_path.replace('.json', '.h5'))
  18. model.compile(loss='binary_crossentropy',
  19. optimizer='adam',
  20. metrics=['accuracy'])
  21. # load scene and its `n` first pixel value data
  22. scene_path = os.path.join(cfg.folder_scenes_path, _scene_name)
  23. for id_column in range(cfg.number_of_columns):
  24. folder_path = os.path.join(scene_path, str(id_column))
  25. pixels_predicted = []
  26. for id_row in range(cfg.number_of_rows):
  27. pixel_filename = _scene_name + '_' + str(id_column) + '_' + str(id_row) + ".dat"
  28. pixel_file_path = os.path.join(folder_path, pixel_filename)
  29. with open(pixel_file_path, 'r') as f:
  30. # predict the expected pixel value
  31. lines = [float(l)/255. for l in f.readlines()]
  32. pixel_values = lines[0:int(_n)]
  33. data = compute_feature(_feature_choice, pixel_values)
  34. pixel_values = np.array(data).reshape(1, (int(_n)))
  35. # predict pixel per pixel
  36. pixels_predicted.append(model.predict(pixel_values))
  37. # change normalized predicted value to pixel value
  38. pixels_predicted = [ val * 255. for val in pixels_predicted]
  39. for id_pixel, pixel in enumerate(pixels_predicted):
  40. output_image[id_pixel, id_column] = pixel
  41. print("{0:.2f}%".format(id_column / cfg.number_of_columns * 100))
  42. sys.stdout.write("\033[F")
  43. return output_image
  44. def main():
  45. parser = argparse.ArgumentParser(description="Train model and saved it")
  46. parser.add_argument('--scene', type=str, help='Scene name to reconstruct', choices=cfg.scenes_list)
  47. parser.add_argument('--model_path', type=str, help='Json model file path')
  48. parser.add_argument('--n', type=str, help='Number of pixel values approximated to keep')
  49. parser.add_argument('--feature', type=str, help='Feature choice to compute from samples', choices=cfg.features_list)
  50. parser.add_argument('--image_name', type=str, help="The ouput image name")
  51. args = parser.parse_args()
  52. param_scene_name = args.scene
  53. param_n = args.n
  54. param_feature = args.feature
  55. param_model_path = args.model_path
  56. param_image_name = args.image_name
  57. # get default value of `n` param
  58. if not param_n:
  59. param_n = param_model_path.split('_')[0]
  60. output_image = reconstruct(param_scene_name, param_model_path, param_n, param_feature)
  61. if not os.path.exists(cfg.reconstructed_folder):
  62. os.makedirs(cfg.reconstructed_folder)
  63. image_path = os.path.join(cfg.reconstructed_folder, param_image_name)
  64. img = Image.fromarray(np.uint8(output_image))
  65. img.save(image_path)
  66. if __name__== "__main__":
  67. main()