generate_dataset.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. """
  4. Created on Wed Jun 19 11:47:42 2019
  5. @author: jbuisine
  6. """
  7. import sys, os, argparse
  8. import numpy as np
  9. import random
  10. import time
  11. import json
  12. from PIL import Image
  13. from ipfml.processing.segmentation import divide_in_blocks
  14. from skimage import color
  15. from modules.utils import config as cfg
  16. from modules.utils import data as dt
  17. from modules.classes.Transformation import Transformation
  18. # getting configuration information
  19. config_filename = cfg.config_filename
  20. zone_folder = cfg.zone_folder
  21. learned_folder = cfg.learned_zones_folder
  22. min_max_filename = cfg.min_max_filename_extension
  23. # define all scenes values
  24. scenes_list = cfg.scenes_names
  25. scenes_indexes = cfg.scenes_indices
  26. choices = cfg.normalization_choices
  27. dataset_path = cfg.dataset_path
  28. zones = cfg.zones_indices
  29. seuil_expe_filename = cfg.seuil_expe_filename
  30. metric_choices = cfg.metric_choices_labels
  31. output_data_folder = cfg.output_data_folder
  32. generic_output_file_svd = '_random.csv'
  33. def generate_data_model(_scenes_list, _filename, _transformations, _scenes, _nb_zones = 4, _random=0):
  34. output_train_filename = _filename + ".train"
  35. output_test_filename = _filename + ".test"
  36. if not '/' in output_train_filename:
  37. raise Exception("Please select filename with directory path to save data. Example : data/dataset")
  38. # create path if not exists
  39. if not os.path.exists(output_data_folder):
  40. os.makedirs(output_data_folder)
  41. zones_indices = zones
  42. train_file_data = []
  43. test_file_data = []
  44. scenes = os.listdir(dataset_path)
  45. # remove min max file from scenes folder
  46. scenes = [s for s in scenes if min_max_filename not in s]
  47. # go ahead each scenes
  48. for folder_scene in _scenes_list:
  49. scene_path = os.path.join(dataset_path, folder_scene)
  50. # shuffle list of zones (=> randomly choose zones)
  51. # only in random mode
  52. if _random:
  53. random.shuffle(zones_indices)
  54. # store zones learned
  55. learned_zones_indices = zones_indices[:_nb_zones]
  56. # write into file
  57. folder_learned_path = os.path.join(learned_folder, _filename.split('/')[1])
  58. if not os.path.exists(folder_learned_path):
  59. os.makedirs(folder_learned_path)
  60. file_learned_path = os.path.join(folder_learned_path, folder_scene + '.csv')
  61. with open(file_learned_path, 'w') as f:
  62. for i in learned_zones_indices:
  63. f.write(str(i) + ';')
  64. for id_zone, index_folder in enumerate(zones_indices):
  65. index_str = str(index_folder)
  66. if len(index_str) < 2:
  67. index_str = "0" + index_str
  68. current_zone_folder = "zone" + index_str
  69. zone_path = os.path.join(scene_path, current_zone_folder)
  70. # custom path for interval of reconstruction and metric
  71. metrics_path = []
  72. for transformation in _transformations:
  73. # check if it's a static content and create augmented images if necessary
  74. if transformation.getName() == 'static':
  75. # {sceneName}/zoneXX/static
  76. static_metric_path = os.path.join(zone_path, transformation.getName())
  77. # img.png
  78. image_name = transformation.getParam().split('/')[-1]
  79. # {sceneName}/zoneXX/static/img
  80. image_folder = image_name.replace('.png', '')
  81. image_folder_path = os.path.join(static_metric_path, image_folder)
  82. if not os.path.exists(image_folder_path):
  83. os.makedirs(image_folder_path)
  84. metrics_path.append(image_folder_path)
  85. # get image path to manage
  86. # {sceneName}/static/img.png
  87. transform_image_path = os.path.join(scene_path, transformation.getName(), image_name)
  88. static_transform_image = Image.open(transform_image_path)
  89. static_transform_image_block = divide_in_blocks(static_transform_image, cfg.keras_img_size)[id_zone]
  90. # compute augmented images if necessary
  91. rotations = [0, 90, 180, 270]
  92. img_flip_labels = ['original', 'horizontal', 'vertical', 'both']
  93. horizontal_img = static_transform_image_block.transpose(Image.FLIP_LEFT_RIGHT)
  94. vertical_img = static_transform_image_block.transpose(Image.FLIP_TOP_BOTTOM)
  95. both_img = static_transform_image_block.transpose(Image.TRANSPOSE)
  96. flip_images = [static_transform_image_block, horizontal_img, vertical_img, both_img]
  97. # rotate and flip image to increase dataset size
  98. for id, flip in enumerate(flip_images):
  99. for rotation in rotations:
  100. rotated_output_img = flip.rotate(rotation)
  101. output_reconstructed_filename = image_folder + cfg.post_image_name_separator
  102. output_reconstructed_filename = output_reconstructed_filename + img_flip_labels[id] + '_' + str(rotation) + '.png'
  103. output_reconstructed_path = os.path.join(image_folder_path, output_reconstructed_filename)
  104. if not os.path.exists(output_reconstructed_path):
  105. rotated_output_img.save(output_reconstructed_path)
  106. else:
  107. metric_interval_path = os.path.join(zone_path, transformation.getTransformationPath())
  108. metrics_path.append(metric_interval_path)
  109. # as labels are same for each metric
  110. for label in os.listdir(metrics_path[0]):
  111. label_metrics_path = []
  112. for path in metrics_path:
  113. label_path = os.path.join(path, label)
  114. label_metrics_path.append(label_path)
  115. # getting images list for each metric
  116. metrics_images_list = []
  117. for index_metric, label_path in enumerate(label_metrics_path):
  118. if _transformations[index_metric].getName() == 'static':
  119. # by default append nothing..
  120. metrics_images_list.append([])
  121. else:
  122. images = sorted(os.listdir(label_path))
  123. metrics_images_list.append(images)
  124. # construct each line using all images path of each
  125. for index_image in range(0, len(metrics_images_list[0])):
  126. images_path = []
  127. # get information about rotation and flip from first transformation (need to be a not static transformation)
  128. current_post_fix = metrics_images_list[0][index_image].split(cfg.post_image_name_separator)[-1]
  129. # getting images with same index and hence name for each metric (transformation)
  130. for index_metric in range(0, len(metrics_path)):
  131. # custom behavior for static transformation (need to check specific image)
  132. if _transformations[index_metric].getName() == 'static':
  133. # add static path with selecting correct data augmented image
  134. image_name = _transformations[index_metric].getParam().split('/')[-1].replace('.png', '')
  135. img_path = os.path.join(metrics_path[index_metric], image_name + cfg.post_image_name_separator + current_post_fix)
  136. images_path.append(img_path)
  137. else:
  138. img_path = metrics_images_list[index_metric][index_image]
  139. images_path.append(os.path.join(label_metrics_path[index_metric], img_path))
  140. if label == cfg.noisy_folder:
  141. line = '1;'
  142. else:
  143. line = '0;'
  144. # compute line information with all images paths
  145. for id_path, img_path in enumerate(images_path):
  146. if id_path < len(images_path) - 1:
  147. line = line + img_path + '::'
  148. else:
  149. line = line + img_path
  150. line = line + '\n'
  151. if id_zone < _nb_zones and folder_scene in _scenes:
  152. train_file_data.append(line)
  153. else:
  154. test_file_data.append(line)
  155. train_file = open(output_train_filename, 'w')
  156. test_file = open(output_test_filename, 'w')
  157. random.shuffle(train_file_data)
  158. random.shuffle(test_file_data)
  159. for line in train_file_data:
  160. train_file.write(line)
  161. for line in test_file_data:
  162. test_file.write(line)
  163. train_file.close()
  164. test_file.close()
  165. def main():
  166. parser = argparse.ArgumentParser(description="Compute specific dataset for model using of metric")
  167. parser.add_argument('--output', type=str, help='output file name desired (.train and .test)')
  168. parser.add_argument('--metrics', type=str,
  169. help="list of metrics choice in order to compute data",
  170. default='svd_reconstruction, ipca_reconstruction',
  171. required=True)
  172. parser.add_argument('--params', type=str,
  173. help="list of specific param for each metric choice (See README.md for further information in 3D mode)",
  174. default='100, 200 :: 50, 25',
  175. required=True)
  176. parser.add_argument('--scenes', type=str, help='List of scenes to use for training data')
  177. parser.add_argument('--nb_zones', type=int, help='Number of zones to use for training data set', choices=list(range(1, 17)))
  178. parser.add_argument('--renderer', type=str, help='Renderer choice in order to limit scenes used', choices=cfg.renderer_choices, default='all')
  179. parser.add_argument('--random', type=int, help='Data will be randomly filled or not', choices=[0, 1])
  180. args = parser.parse_args()
  181. p_filename = args.output
  182. p_metrics = list(map(str.strip, args.metrics.split(',')))
  183. p_params = list(map(str.strip, args.params.split('::')))
  184. p_scenes = args.scenes.split(',')
  185. p_nb_zones = args.nb_zones
  186. p_renderer = args.renderer
  187. p_random = args.random
  188. # create list of Transformation
  189. transformations = []
  190. for id, metric in enumerate(p_metrics):
  191. if metric not in metric_choices:
  192. raise ValueError("Unknown metric, please select a correct metric : ", metric_choices)
  193. transformations.append(Transformation(metric, p_params[id]))
  194. if transformations[0].getName() == 'static':
  195. raise ValueError("The first transformation in list cannot be static")
  196. # list all possibles choices of renderer
  197. scenes_list = dt.get_renderer_scenes_names(p_renderer)
  198. scenes_indices = dt.get_renderer_scenes_indices(p_renderer)
  199. # getting scenes from indexes user selection
  200. scenes_selected = []
  201. for scene_id in p_scenes:
  202. index = scenes_indices.index(scene_id.strip())
  203. scenes_selected.append(scenes_list[index])
  204. # create database using img folder (generate first time only)
  205. generate_data_model(scenes_list, p_filename, transformations, scenes_selected, p_nb_zones, p_random)
  206. if __name__== "__main__":
  207. main()