splitTrainAndTest.py 673 B

1234567891011121314151617181920212223242526272829
  1. import random
  2. import os
  3. import subprocess
  4. import sys
  5. def split_data_set(image_dir):
  6. f_val = open("snowman_test.txt", 'w')
  7. f_train = open("snowman_train.txt", 'w')
  8. path, dirs, files = next(os.walk(image_dir))
  9. data_size = len(files)
  10. ind = 0
  11. data_test_size = int(0.1 * data_size)
  12. test_array = random.sample(range(data_size), k=data_test_size)
  13. for f in os.listdir(image_dir):
  14. if(f.split(".")[1] == "jpg"):
  15. ind += 1
  16. if ind in test_array:
  17. f_val.write(image_dir+'/'+f+'\n')
  18. else:
  19. f_train.write(image_dir+'/'+f+'\n')
  20. split_data_set(sys.argv[1])