reconstruct_scene_mean.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. import numpy as np
  2. import pandas as pd
  3. import os, sys, argparse
  4. from sklearn import linear_model
  5. from sklearn import svm
  6. from sklearn.utils import shuffle
  7. import modules.config as cfg
  8. import modules.metrics as metrics
  9. from joblib import dump, load
  10. from PIL import Image
  11. def reconstruct(_scene_name, _n):
  12. # construct the empty output image
  13. output_image = np.empty([cfg.number_of_rows, cfg.number_of_columns])
  14. # load scene and its `n` first pixel value data
  15. scene_path = os.path.join(cfg.folder_scenes_path, _scene_name)
  16. for id_column in range(cfg.number_of_columns):
  17. folder_path = os.path.join(scene_path, str(id_column))
  18. for id_row in range(cfg.number_of_rows):
  19. pixel_filename = _scene_name + '_' + str(id_column) + '_' + str(id_row) + ".dat"
  20. pixel_file_path = os.path.join(folder_path, pixel_filename)
  21. with open(pixel_file_path, 'r') as f:
  22. # predict the expected pixel value
  23. lines = [float(l) for l in f.readlines()]
  24. mean = sum(lines[0:int(_n)]) / float(_n)
  25. output_image[id_row, id_column] = mean
  26. print("{0:.2f}%".format(id_column / cfg.number_of_columns * 100))
  27. sys.stdout.write("\033[F")
  28. return output_image
  29. def main():
  30. parser = argparse.ArgumentParser(description="Train model and saved it")
  31. parser.add_argument('--scene', type=str, help='Scene name to reconstruct', choices=cfg.scenes_list)
  32. parser.add_argument('--n', type=str, help='Number of samples to take')
  33. parser.add_argument('--image_name', type=str, help="The ouput image name")
  34. args = parser.parse_args()
  35. param_scene_name = args.scene
  36. param_n = args.n
  37. param_image_name = args.image_name
  38. output_image = reconstruct(param_scene_name, param_n)
  39. if not os.path.exists(cfg.reconstructed_folder):
  40. os.makedirs(cfg.reconstructed_folder)
  41. image_path = os.path.join(cfg.reconstructed_folder, param_image_name)
  42. img = Image.fromarray(np.uint8(output_image))
  43. img.save(image_path)
  44. if __name__== "__main__":
  45. main()