File size: 822 Bytes
45928a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import os
import pandas as pd
from sklearn.model_selection import train_test_split
DATA_DIR = "data/train"
CSV_PATH = "data/trainLabels.csv"
TEST_SIZE = 0.2
RANDOM_STATE = 42
# Load the CSV file into a pandas DataFrame and add the image path
df = pd.read_csv(CSV_PATH, names=['image_path', 'label'], converters={'image_path': lambda x: f"{DATA_DIR}/{x}.jpeg"})
# drop row where image does not exist
df = df[df['image_path'].apply(lambda x: os.path.exists(x))]
# split the data into train and validation sets such that the class distribution is the same in both sets
df_train, df_val = train_test_split(df, test_size=TEST_SIZE, stratify=df['label'], random_state=RANDOM_STATE)
# Save the train and validation sets to CSV files
df_train.to_csv("data/train.csv", index=False)
df_val.to_csv("data/val.csv", index=False)
|