display_reconstructed_image_from_humans.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. # main imports
  2. import numpy as np
  3. import pandas as pd
  4. import math
  5. import time
  6. import os, sys, argparse
  7. # image processing imports
  8. import matplotlib.pyplot as plt
  9. from PIL import Image
  10. # modules imports
  11. sys.path.insert(0, '') # trick to enable import of main folder module
  12. import custom_config as cfg
  13. from data_attributes import get_image_features
  14. from modules.utils import data as dt
  15. # other variables
  16. learned_zones_folder = cfg.learned_zones_folder
  17. models_name = cfg.models_names_list
  18. # utils information
  19. zone_width, zone_height = (200, 200)
  20. scene_width, scene_height = (800, 800)
  21. nb_x_parts = math.floor(scene_width / zone_width)
  22. def reconstruct_image(scene_name, output):
  23. """
  24. @brief Method used to display simulation given .csv files
  25. @param scene_name, scene name used
  26. @param output, the output filename
  27. @return nothing
  28. """
  29. # compute zone start index
  30. zones_coordinates = []
  31. for zone_index in cfg.zones_indices:
  32. x_zone = (zone_index % nb_x_parts) * zone_width
  33. y_zone = (math.floor(zone_index / nb_x_parts)) * zone_height
  34. zones_coordinates.append((x_zone, y_zone))
  35. scene_folder = os.path.join(cfg.dataset_path, scene_name)
  36. folder_scene_elements = os.listdir(scene_folder)
  37. zones_folder = [zone for zone in folder_scene_elements if 'zone' in zone]
  38. zones_folder = sorted(zones_folder)
  39. scenes_images = [img for img in folder_scene_elements if cfg.scene_image_extension in img]
  40. scenes_images = sorted(scenes_images)
  41. # 1. find thresholds from scene
  42. human_thresholds = []
  43. for zone_folder in zones_folder:
  44. zone_path = os.path.join(scene_folder, zone_folder)
  45. with open(os.path.join(zone_path, cfg.seuil_expe_filename)) as f:
  46. human_thresholds.append(int(f.readline()))
  47. # 2. find images for each zone which are attached to these human thresholds by the model
  48. zone_images_index = []
  49. for threshold in human_thresholds:
  50. current_image_index = 0
  51. for image_name in scenes_images:
  52. image_quality = dt.get_scene_image_quality(image_name)
  53. if image_quality > threshold:
  54. current_image_index = image_quality
  55. break
  56. str_index = str(current_image_index)
  57. while len(str_index) < 5:
  58. str_index = "0" + str_index
  59. zone_images_index.append(str_index)
  60. images_zones = []
  61. line_images_zones = []
  62. # get image using threshold by zone
  63. for id, zone_index in enumerate(zone_images_index):
  64. filtered_images = [img for img in scenes_images if zone_index in img]
  65. if len(filtered_images) > 0:
  66. image_name = filtered_images[0]
  67. else:
  68. image_name = scenes_images[-1]
  69. image_path = os.path.join(scene_folder, image_name)
  70. selected_image = Image.open(image_path)
  71. x_zone, y_zone = zones_coordinates[id]
  72. zone_image = np.array(selected_image)[y_zone:y_zone+zone_height, x_zone:x_zone+zone_width]
  73. line_images_zones.append(zone_image)
  74. if int(id + 1) % int(scene_width / zone_width) == 0:
  75. images_zones.append(np.concatenate(line_images_zones, axis=1))
  76. line_images_zones = []
  77. # 3. reconstructed the image using these zones
  78. reconstructed_image = np.concatenate(images_zones, axis=0)
  79. # 4. Save the image with generated name based on scene
  80. reconstructed_pil_img = Image.fromarray(reconstructed_image)
  81. folders = output.split('/')
  82. if len(folders) > 1:
  83. output_folder = '/'.join(folders[:len(folders) - 1])
  84. if not os.path.exists(output_folder):
  85. os.makedirs(output_folder)
  86. reconstructed_pil_img.save(output)
  87. def main():
  88. parser = argparse.ArgumentParser(description="Compute and save reconstructed images from human thresholds")
  89. parser.add_argument('--scene', type=str, help='Scene index to use', choices=cfg.scenes_indices)
  90. parser.add_argument('--output', type=str, help='Output reconstructed image path and filename')
  91. args = parser.parse_args()
  92. p_scene = args.scene
  93. p_output = args.output
  94. scenes_list = cfg.scenes_names
  95. scenes_indices = cfg.scenes_indices
  96. scene_index = scenes_indices.index(p_scene.strip())
  97. scene_name = scenes_list[scene_index]
  98. reconstruct_image(scene_name, p_output)
  99. if __name__== "__main__":
  100. main()