generate_selected_zones_file.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104
  1. # main imports
  2. import numpy as np
  3. import pandas as pd
  4. import sys, os, argparse
  5. import random
  6. # modules and config imports
  7. sys.path.insert(0, '') # trick to enable import of main folder module
  8. import custom_config as cfg
  9. def save_learned_zones(output_name, scene, zones):
  10. if not os.path.exists(cfg.output_zones_learned):
  11. os.makedirs(cfg.output_zones_learned)
  12. with open(os.path.join(cfg.output_zones_learned, output_name), 'a') as f:
  13. f.write(scene + ';')
  14. for zone in zones:
  15. f.write(str(zone) + ';')
  16. f.write('\n')
  17. def get_random_zones(scene, zones, n_zones):
  18. random.shuffle(zones)
  19. # specific case for 'Cuisine01' (zone 12 is also noisy even in reference image)
  20. # if scene == 'Cuisine01':
  21. # while 12 in zones[0:n_zones]:
  22. # random.shuffle(zones)
  23. return zones[0:n_zones]
  24. def main():
  25. parser = argparse.ArgumentParser(description="Read and compute entropy data file (using diff)")
  26. parser.add_argument('--folder', type=str, help='dataset scene folder', required=True)
  27. parser.add_argument('--n_zones', type=int, help='number of zones used in train', default=10)
  28. parser.add_argument('--output', type=str, help='file with specific training zone', required=True)
  29. parser.add_argument('--thresholds', type=str, help='file with specific thresholds (using only scene from this file', default='')
  30. args = parser.parse_args()
  31. p_folder = args.folder
  32. p_n_zones = args.n_zones
  33. p_output = args.output
  34. p_thresholds = args.thresholds
  35. # extract scenes to use if specified
  36. available_scenes = None
  37. if len(p_thresholds) > 0:
  38. available_scenes = []
  39. with open(p_thresholds) as f:
  40. thresholds_line = f.readlines()
  41. for line in thresholds_line:
  42. data = line.split(';')
  43. del data[-1] # remove unused last element `\n`
  44. current_scene = data[0]
  45. # need to rename `current_name` because we only used part6
  46. # scene_split = current_scene.split('_')
  47. # del scene_split[-1]
  48. # scene_name = '_'.join(scene_split)
  49. available_scenes.append(current_scene)
  50. # specific number of zones (zones indices)
  51. zones = np.arange(16)
  52. # get all scene names
  53. scenes = os.listdir(p_folder)
  54. # create output thresholds directory if necessary
  55. folder, _ = os.path.split(p_output)
  56. if len(folder) > 0:
  57. os.makedirs(folder)
  58. # for each scene we generate random zones choice
  59. for folder_scene in scenes:
  60. if available_scenes is not None:
  61. if folder_scene in available_scenes:
  62. selected_zones = get_random_zones(folder_scene, zones, p_n_zones)
  63. save_learned_zones(p_output, folder_scene, selected_zones)
  64. else:
  65. selected_zones = get_random_zones(folder_scene, zones, p_n_zones)
  66. save_learned_zones(p_output, folder_scene, selected_zones)
  67. if __name__== "__main__":
  68. main()