How to split the custom image dataset into training, testing, and validation folders for Machine Learning Projects

Machine learning can seem complex, but it boils down to three key steps when using pre-trained models: training, validation, and testing.

When creating our machine-learning project with a custom dataset built from scratch! We might take photos ourselves, and then apply all the necessary preprocessing steps. Here’s where data splitting becomes crucial. While we could manually divide our data into training, validation, and testing folders, who wants that tedious task? Luckily, we can automate the process with some handy tools!

There’s no one-size-fits-all approach for split ratios, but common practice is 70% training, 15% validation, and 15% testing. This ensures our model has enough data to learn (training), can be fine-tuned during training (validation), and is ultimately evaluated on unseen data (testing).

Let’s see the code now…

1) Imports

We’ll import (os, random, and shutil) that are built-in modules in Python that provide functionalities for various tasks, such as:

• os to navigate directories and access image files.
• random to randomly select a portion of images for each set.
• shutil to copy those selected images to their respective training, testing, and validation folders.

import os
import shutil
import random

2) Directories’ Path

Give the paths for the input folder, along with the path for the train, test, and validate folders. It is better to create the 3 folders in the desired location beforehand but if not created the code will automatically create them for you.

# GIVE YOUR FOLDER PATHS. given below folder paths are for example only.

#giving the path for each folder.
input_dir = "D:\Project\RawData" 
test_dir = "D:\Project\Train"
train_dir = "D:\Project\Test"
val_dir = "D:\Project\Validate"

3) Set the desired split ratios

train_ratio = 0.7
val_ratio = 0.15
test_ratio = 0.15

4) Create the train, val, and test directories if they don’t already exist

if not os.path.exists(train_dir):
    os.makedirs(train_dir)
if not os.path.exists(val_dir):
    os.makedirs(val_dir)
if not os.path.exists(test_dir):
    os.makedirs(test_dir)

5) Get a list of all the image files in the input directory

Get KIKI’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The code os.listdir(input_dir) returns a list of all files and directories in the specified input_dir directory. The list comprehension [f for f in os.listdir(input_dir) if f.endswith('.jpg')] filters this list to include only the files that end with the .jpg extension, which typically indicates that they are image files. The resulting list, image_files, contains the names of all the image files in the input_dir directory.

image_files = [f for f in os.listdir(input_dir) if f.endswith('.jpg')]

6) Randomize the order of the image files

To randomize the order of the image files in the image_files list, we use the random.shuffle() function from the random module. This will shuffle the elements in the image_files list so that the order of the image files is randomized. The code below will create a new list image_files_shuffled with the same elements as image_files but in a random order.

random.shuffle(image_files)

7) Calculate the indices to split the data into train, val, and test sets

This code calculates the indices to split the image_files list based on the specified train_ratio, val_ratio, and test_ratio. These indices can be used to split the image_files list into train, validation, and test sets using slicing.

train_indices = int(len(image_files) * train_ratio)
val_indices = int(len(image_files) * (train_ratio + val_ratio))
test_indices = int(len(image_files) * (train_ratio + val_ratio + test_ratio))

8) Split the image files into train, val, and test sets

train_files = image_files[:train_indices]
val_files = image_files[train_indices:val_indices]
test_files = image_files[val_indices:]

9) Copy the image files into the train, val, and test directories

for file in train_files:
    shutil.copy(os.path.join(input_dir, file), os.path.join(train_dir, file))
    image_files.remove(file)
for file in val_files:
    shutil.copy(os.path.join(input_dir, file), os.path.join(val_dir, file))
    image_files.remove(file)
for file in test_files:
    shutil.copy(os.path.join(input_dir, file), os.path.join(test_dir, file))
    image_files.remove(file)

Alrighty, folks! Once you run that code, head over to the folder location and take a peek. You should see three shiny new folders just bursting with images, all neatly sorted into their training, validation, and testing teams. High fives all around! We just conquered a super important step in our machine-learning adventure. Now you can dive into the fun stuff — annotating those training set images and tackling the next steps.

P.S. If this code made your life easier, don’t be shy, give it a big ol’ virtual clap!!!👏

Thanks for reading! 😃

Below is the whole code given, Don’t forget to change the folder paths to your own! :

import os
import shutil
import random

# GIVE PATHS OF YOUR  
#giving the path for each folder.
input_dir = "D:\\...\\RawData" 
test_dir = "D:\\...\\Train"
train_dir = "D:\\...\\Test"
val_dir = "D:\\...\\Validate"


# Set the desired split ratios
train_ratio = 0.7
val_ratio = 0.15
test_ratio = 0.15

# Create the train, val, and test directories if they don't already exist
if not os.path.exists(train_dir):
    os.makedirs(train_dir)
if not os.path.exists(val_dir):
    os.makedirs(val_dir)
if not os.path.exists(test_dir):
    os.makedirs(test_dir)

# Get a list of all the image files in the input directory
# if your images are of .png or .jpeg then use it instead of .jpg
image_files = [f for f in os.listdir(input_dir) if f.endswith('.jpg')]

# Randomize the order of the image files
random.shuffle(image_files)

# Calculate the indices to split the data into train, val, and test sets
train_indices = int(len(image_files) * train_ratio)
val_indices = int(len(image_files) * (train_ratio + val_ratio))
test_indices = int(len(image_files) * (train_ratio + val_ratio + test_ratio))

# Split the image files into train, val, and test sets
train_files = image_files[:train_indices]
val_files = image_files[train_indices:val_indices]
test_files = image_files[val_indices:]

# Copy the image files into the train, val, and test directories
for file in train_files:
    shutil.copy(os.path.join(input_dir, file), os.path.join(train_dir, file))
    image_files.remove(file)
for file in val_files:
    shutil.copy(os.path.join(input_dir, file), os.path.join(val_dir, file))
    image_files.remove(file)
for file in test_files:
    shutil.copy(os.path.join(input_dir, file), os.path.join(test_dir, file))
    image_files.remove(file)

In Plain English 🚀

Thank you for being a part of the In Plain English community! Before you go:

• Be sure to clap and follow the writer ️👏️️
• Follow us: X | LinkedIn | YouTube | Discord | Newsletter
• Visit our other platforms: Stackademic | CoFeed | Venture | Cubed
• More content at PlainEnglish.io