How to do train-test split for multi-class classification

Question

I am performing multi-class classification problem of different concentrations of Acetaminophen in a specified dataset. My data is in the form of images and I am using CNN. I have compiled all the image paths (with the images being the features) and their respective concentrations (classes) into one excel. Next, I formed a variable "X" to contain all the dataset image arrays and y to contain all the labels, as shown in the code below:

img_path = pd.read_excel(os.path.join(Dir,"DPV Images.xlsx"))

X = []
y = []

for index, row in img_path.iterrows():
    img_array = cv2.imread(row[0]) # Reads the image
    img_class = row[1] # Contains the label
    X.append(img_array)
    y.append(img_class)

My question is: How can I do a train-test split in this case, to ensure that the train and test data is split evenly between all classes. For example, using scikitlearn's train_test_split with test size of 0.2 will divide the entire dataset into 80% training and 20% training, but it does not ensure that each class is divided into 80% training and 20% testing. How can I go about this?

Answer 1

To do the train-test split in a method that assures an equal distribution of classes between the training and testing sets, utilize the StratifiedShuffleSplit class fromscikit-learn model selection module.

Try:

from sklearn.model_selection import StratifiedShuffleSplit
splitter = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=0)
for train_idx, test_idx in splitter.split(X, y):
    X_train = [X[i] for i in train_idx]
    y_train = [y[i] for i in train_idx]
    X_test = [X[i] for i in test_idx]
    y_test = [y[i] for i in test_idx]

Answer 2

You can write a custom function to do this job, below is an example:

import numpy as np
from sklearn.model_selection import train_test_split

def split_data(X, test_size=0.25, random_state=2):
    y=X["Label"]
    unique_classes = np.unique(y)
    X_train = []
    X_test = []
    y_train = []
    y_test = []
    for class_value in unique_classes:
        X_class = X[y == class_value]
        y_class = y[y == class_value]
        X_class_train, X_class_test, y_class_train, y_class_test = train_test_split(X_class, y_class, test_size=test_size, random_state=random_state)
        X_train.append(X_class_train)
        X_test.append(X_class_test)
        y_train.append(y_class_train)
        y_test.append(y_class_test)
    X_train = np.concatenate(X_train)
    X_test = np.concatenate(X_test)
    y_train = np.concatenate(y_train)
    y_test = np.concatenate(y_test)
    return X_train, X_test, y_train, y_test