Make script work with nested folders

Author: EdjeElectronics

Diff for: util_scripts/train_val_test_split.py

@@ -7,6 +7,7 @@
 # https://colab.research.google.com/github/EdjeElectronics/TensorFlow-Lite-Object-Detection-on-Android-and-Raspberry-Pi/blob/master/Train_TFLite2_Object_Detction_Model.ipynb
 
 import glob
+from pathlib import Path
 import random
 import os
 
@@ -17,10 +18,10 @@
 test_path = '/content/images/test'
 
 # Get list of all images
-jpg_file_list = glob.glob(image_path + '/*.jpg')
-JPG_file_list = glob.glob(image_path + '/*.JPG')
-png_file_list = glob.glob(image_path + '/*.png')
-bmp_file_list = glob.glob(image_path + '/*.bmp')
+jpg_file_list = [path for path in Path(image_path).rglob('*.jpg')]
+JPG_file_list = [path for path in Path(image_path).rglob('*.JPG')]
+png_file_list = [path for path in Path(image_path).rglob('*.png')]
+bmp_file_list = [path for path in Path(image_path).rglob('*.bmp')]
 
 file_list = jpg_file_list + JPG_file_list + png_file_list + bmp_file_list
 file_num = len(file_list)
@@ -33,35 +34,39 @@
 train_num = int(file_num*train_percent)
 val_num = int(file_num*val_percent)
 test_num = file_num - train_num - val_num
-print('Files moving to train: %d' % train_num)
-print('Files moving to validation: %d' % val_num)
-print('Files moving to test: %d' % test_num)
+print('Images moving to train: %d' % train_num)
+print('Images moving to validation: %d' % val_num)
+print('Images moving to test: %d' % test_num)
 
 # Select 80% of files randomly and move them to train folder
 for i in range(train_num):
     move_me = random.choice(file_list)
-    fn = move_me.split('/')[-1]
-    base_fn = fn[:(len(fn)-4)] # Gets rid of .jpg, .png, or .bmp at end of the string
+    fn = move_me.name
+    base_fn = move_me.stem
+    parent_path = move_me.parent
     xml_fn = base_fn + '.xml'
     os.rename(move_me, train_path+'/'+fn)
-    os.rename(os.path.join(image_path,xml_fn),os.path.join(train_path,xml_fn))
+    os.rename(os.path.join(parent_path,xml_fn),os.path.join(train_path,xml_fn))
     file_list.remove(move_me)
 
 # Select 10% of remaining files and move them to validation folder
 for i in range(val_num):
     move_me = random.choice(file_list)
-    fn = move_me.split('/')[-1]
-    base_fn = fn[:(len(fn)-4)] # Gets rid of .jpg, .png, or .bmp at end of the string
+    fn = move_me.name
+    base_fn = move_me.stem
+    parent_path = move_me.parent
     xml_fn = base_fn + '.xml'
     os.rename(move_me, val_path+'/'+fn)
-    os.rename(os.path.join(image_path,xml_fn),os.path.join(val_path,xml_fn))
+    os.rename(os.path.join(parent_path,xml_fn),os.path.join(val_path,xml_fn))
     file_list.remove(move_me)
 
 # Move remaining files to test folder
 for i in range(test_num):
-    file = file_list[i]
-    fn = file.split('/')[-1]
-    base_fn = fn[:(len(fn)-4)] # Gets rid of .jpg, .png, or .bmp at end of the string
+    move_me = random.choice(file_list)
+    fn = move_me.name
+    base_fn = move_me.stem
+    parent_path = move_me.parent
     xml_fn = base_fn + '.xml'
-    os.rename(file, test_path+'/'+fn)
-    os.rename(os.path.join(image_path,xml_fn),os.path.join(test_path,xml_fn))
+    os.rename(move_me, test_path+'/'+fn)
+    os.rename(os.path.join(parent_path,xml_fn),os.path.join(test_path,xml_fn))
+    file_list.remove(move_me)