generate_dataset.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. """
  4. Created on Fri Sep 14 21:02:42 2018
  5. @author: jbuisine
  6. """
  7. from __future__ import print_function
  8. import glob, image_slicer
  9. import sys, os, getopt
  10. from PIL import Image
  11. import shutil
  12. # show to create own dataset https://gist.github.com/fchollet/0830affa1f7f19fd47b06d4cf89ed44d
  13. NUMBER_SUB_IMAGES = 100
  14. def init_directory():
  15. if os.path.exists('data'):
  16. print("Removing all previous data...")
  17. shutil.rmtree('data')
  18. if not os.path.exists('data'):
  19. print("Creating new data...")
  20. os.makedirs('data/train/final')
  21. os.makedirs('data/train/noisy')
  22. os.makedirs('data/validation/final')
  23. os.makedirs('data/validation/noisy')
  24. def create_images(folder, output_folder):
  25. images_path = glob.glob(folder + "/*.png")
  26. for img in images_path:
  27. image_name = img.replace(folder, '').replace('/', '')
  28. tiles = image_slicer.slice(img, NUMBER_SUB_IMAGES, save = False)
  29. image_slicer.save_tiles(tiles, directory=output_folder, prefix='part_'+image_name)
  30. def generate_dataset():
  31. create_images('img_train/final', 'data/train/final')
  32. create_images('img_train/noisy', 'data/train/noisy')
  33. create_images('img_validation/final', 'data/validation/final')
  34. create_images('img_validation/noisy', 'data/validation/noisy')
  35. def main():
  36. global NUMBER_SUB_IMAGES
  37. if len(sys.argv) <= 1:
  38. print('Please specify nb sub image per image parameter (use -h if you want to know values)...')
  39. print('generate_dataset.py --nb xxxx')
  40. sys.exit(2)
  41. try:
  42. opts, args = getopt.getopt(sys.argv[1:], "h:n", ["help", "nb="])
  43. except getopt.GetoptError:
  44. # print help information and exit:
  45. print('generate_dataset.py --nb xxxx')
  46. sys.exit(2)
  47. for o, a in opts:
  48. if o == "--help":
  49. print('generate_dataset.py --nb xxxx')
  50. print('20x20 : 1600')
  51. print('40x40 : 400')
  52. print('60x60 : 178 (approximately)')
  53. print('80x80 : 100')
  54. print('100x100 : 64')
  55. sys.exit()
  56. elif o == '--nb':
  57. NUMBER_SUB_IMAGES = int(a)
  58. init_directory()
  59. # create database using img folder (generate first time only)
  60. generate_dataset()
  61. if __name__== "__main__":
  62. main()