estimate_thresholds_cnn.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  1. # main imports
  2. import numpy as np
  3. import pandas as pd
  4. import sys, os, argparse
  5. # image processing
  6. from PIL import Image
  7. from ipfml import utils
  8. from ipfml.processing import transform, segmentation
  9. import matplotlib.pyplot as plt
  10. # model imports
  11. import joblib
  12. from keras.models import load_model
  13. from keras import backend as K
  14. # modules and config imports
  15. sys.path.insert(0, '') # trick to enable import of main folder module
  16. import custom_config as cfg
  17. import modules.utils.data as dt
  18. from modules.classes.Transformation import Transformation
  19. def write_progress(progress):
  20. barWidth = 180
  21. output_str = "["
  22. pos = barWidth * progress
  23. for i in range(barWidth):
  24. if i < pos:
  25. output_str = output_str + "="
  26. elif i == pos:
  27. output_str = output_str + ">"
  28. else:
  29. output_str = output_str + " "
  30. output_str = output_str + "] " + str(int(progress * 100.0)) + " %\r"
  31. print(output_str)
  32. sys.stdout.write("\033[F")
  33. def main():
  34. parser = argparse.ArgumentParser(description="Read and compute entropy data file")
  35. parser.add_argument('--model', type=str, help='model .h5 file')
  36. parser.add_argument('--folder', type=str,
  37. help='folder where scene dataset is available',
  38. required=True)
  39. parser.add_argument('--features', type=str,
  40. help="list of features choice in order to compute data",
  41. default='svd_reconstruction, ipca_reconstruction',
  42. required=True)
  43. parser.add_argument('--params', type=str,
  44. help="list of specific param for each feature choice (See README.md for further information in 3D mode)",
  45. default='100, 200 :: 50, 25',
  46. required=True)
  47. parser.add_argument('--size', type=str,
  48. help="specific size of image",
  49. default='100, 100',
  50. required=True)
  51. parser.add_argument('--n_stop', type=int, help='number of detection to make sure to stop', default=1)
  52. parser.add_argument('--save', type=str, help='filename where to save input data')
  53. parser.add_argument('--label', type=str, help='label to use when saving thresholds')
  54. args = parser.parse_args()
  55. p_model = args.model
  56. p_folder = args.folder
  57. p_features = list(map(str.strip, args.features.split(',')))
  58. p_params = list(map(str.strip, args.params.split('::')))
  59. p_size = args.size
  60. p_n_stop = args.n_stop
  61. p_save = args.save
  62. p_label = args.label
  63. # 1. Load expected transformations
  64. # list of transformations
  65. transformations = []
  66. for id, feature in enumerate(p_features):
  67. if feature not in cfg.features_choices_labels or feature == 'static':
  68. raise ValueError("Unknown feature, please select a correct feature (`static` excluded) : ", cfg.features_choices_labels)
  69. transformations.append(Transformation(feature, p_params[id], p_size))
  70. # 2. load model and compile it
  71. # TODO : check kind of model
  72. model = load_model(p_model)
  73. # model.compile(loss='binary_crossentropy',
  74. # optimizer='rmsprop',
  75. # metrics=['accuracy'])
  76. estimated_thresholds = []
  77. n_estimated_thresholds = []
  78. scene_path = p_folder
  79. if not os.path.exists(scene_path):
  80. print('Unvalid scene path:', scene_path)
  81. exit(0)
  82. # 3. retrieve human_thresholds
  83. # construct zones folder
  84. zones_indices = np.arange(16)
  85. zones_list = []
  86. for index in zones_indices:
  87. index_str = str(index)
  88. while len(index_str) < 2:
  89. index_str = "0" + index_str
  90. zones_list.append(cfg.zone_folder + index_str)
  91. # 4. get estimated thresholds using model and specific method
  92. images_path = sorted([os.path.join(scene_path, img) for img in os.listdir(scene_path) if cfg.scene_image_extension in img])
  93. number_of_images = len(images_path)
  94. image_indices = [ dt.get_scene_image_quality(img_path) for img_path in images_path ]
  95. image_counter = 0
  96. # append empty list
  97. for _ in zones_list:
  98. estimated_thresholds.append(None)
  99. n_estimated_thresholds.append(0)
  100. for img_i, img_path in enumerate(images_path):
  101. blocks = segmentation.divide_in_blocks(Image.open(img_path), (200, 200))
  102. for index, block in enumerate(blocks):
  103. if estimated_thresholds[index] is None:
  104. transformed_list = []
  105. # compute data here
  106. for transformation in transformations:
  107. transformed = transformation.getTransformedImage(block)
  108. transformed_list.append(transformed)
  109. data = np.array(transformed_list)
  110. # compute input size
  111. n_chanels, _, _ = data.shape
  112. if K.image_data_format() == 'chanels_first':
  113. if n_chanels > 1:
  114. data = np.expand_dims(data, axis=0)
  115. else:
  116. if n_chanels > 1:
  117. data = data.transpose()
  118. data = np.expand_dims(data, axis=0)
  119. else:
  120. data = data.transpose()
  121. data = np.expand_dims(data, axis=0)
  122. probs = model.predict(np.array(data))[0]
  123. prediction = list(probs).index(max(probs))
  124. #print(index, ':', image_indices[img_i], '=>', prediction)
  125. if prediction == 0:
  126. n_estimated_thresholds[index] += 1
  127. # if same number of detection is attempted
  128. if n_estimated_thresholds[index] >= p_n_stop:
  129. estimated_thresholds[index] = image_indices[img_i]
  130. else:
  131. n_estimated_thresholds[index] = 0
  132. # write progress bar
  133. write_progress((image_counter + 1) / number_of_images)
  134. image_counter = image_counter + 1
  135. # default label
  136. for i, _ in enumerate(zones_list):
  137. if estimated_thresholds[i] == None:
  138. estimated_thresholds[i] = image_indices[-1]
  139. # 6. save estimated thresholds into specific file
  140. print('\nEstimated thresholds', estimated_thresholds)
  141. if p_save is not None:
  142. with open(p_save, 'a') as f:
  143. f.write(p_label + ';')
  144. for t in estimated_thresholds:
  145. f.write(str(t) + ';')
  146. f.write('\n')
  147. if __name__== "__main__":
  148. main()