generate_dataset.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. #!/usr/bin/env python2
  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 os, glob, image_slicer
  9. from PIL import Image
  10. # show to create own dataset https://gist.github.com/fchollet/0830affa1f7f19fd47b06d4cf89ed44d
  11. NUMBER_SUB_IMAGES = 100
  12. def init_directory():
  13. if not os.path.exists('data'):
  14. os.makedirs('data/train/final')
  15. os.makedirs('data/train/noisy')
  16. os.makedirs('data/validation/final')
  17. os.makedirs('data/validation/noisy')
  18. def create_images(folder, output_folder):
  19. images_path = glob.glob(folder + "/*.png")
  20. for img in images_path:
  21. image_name = img.replace(folder, '').replace('/', '')
  22. tiles = image_slicer.slice(img, NUMBER_SUB_IMAGES, save = False)
  23. image_slicer.save_tiles(tiles, directory=output_folder, prefix='part_'+image_name)
  24. def generate_dataset():
  25. create_images('img_train/final', 'data/train/final')
  26. create_images('img_train/noisy', 'data/train/noisy')
  27. create_images('img_validation/final', 'data/validation/final')
  28. create_images('img_validation/noisy', 'data/validation/noisy')
  29. def main():
  30. init_directory()
  31. # create database using img folder (generate first time only)
  32. generate_dataset()
  33. if __name__== "__main__":
  34. main()