ability to split out train/test/valid folders on export #360
garminTestPilot50
started this conversation in
Ideas
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
-
For yolo training, it is super useful to be able to export out the train/test/valid folders and be able to control the mix (like 80% / 10% / 10%) for the images, and randomly assign the images & labels to those folders with proper pairing.
I ended up writing some python code to move the images & labels afterwards, but it would be great to be able to do that from the tool itself.
Here is the simple code if it helps:
import glob
import random
import os
import shutil
Use raw strings for Windows file paths
PATH = r'F:\my\folder\data\'
img_paths = glob.glob(PATH + r'images*.jpg')
txt_paths = glob.glob(PATH + r'labels*.txt')
Calculate number of files for training, validation
data_size = len(img_paths)
train_size = int(data_size * 0.8)
test_size = int(data_size * 0.1)
valid_size = data_size - train_size - test_size # Ensure all data is used
Shuffle two lists
img_txt = list(zip(img_paths, txt_paths))
random.seed(43)
random.shuffle(img_txt)
img_paths, txt_paths = zip(*img_txt)
Now split them
train_img_paths = img_paths[:train_size]
train_txt_paths = txt_paths[:train_size]
test_img_paths = img_paths[train_size:train_size + test_size]
test_txt_paths = txt_paths[train_size:train_size + test_size]
valid_img_paths = img_paths[train_size + test_size:]
valid_txt_paths = txt_paths[train_size + test_size:]
Function to create folders and move files
def create_and_move_files(folder_path, img_paths, txt_paths):
if not os.path.exists(folder_path):
os.makedirs(folder_path, exist_ok=True) # Use makedirs to create intermediate directories if needed
for img_path, txt_path in zip(img_paths, txt_paths):
try:
shutil.move(img_path, folder_path)
shutil.move(txt_path, folder_path)
except Exception as e:
print(f"An error occurred while moving files: {e}")
Move them to train, valid, test folders
create_and_move_files(PATH + r'train\', train_img_paths, train_txt_paths)
create_and_move_files(PATH + r'test\', test_img_paths, test_txt_paths)
create_and_move_files(PATH + r'valid\', valid_img_paths, valid_txt_paths)
Beta Was this translation helpful? Give feedback.
All reactions