EmanHussein/ResNet50_Garbage_Classifier

1import os2import shutil3from sklearn.model_selection import train_test_split4 5# Set paths6try:7    ORIGINAL_DATASET = input("Enter the path to the original dataset: ")  # Path to the original dataset8    if not ORIGINAL_DATASET:9        raise ValueError("Please provide the path to the original dataset.")10    BASE_DIR = "../data/preprocessed_data"  # Output directory11 12    if not os.path.exists(BASE_DIR):13        os.makedirs(BASE_DIR)14except Exception as e:15    print(f"Error: {e}")16    exit(1)17 18# Check if the original dataset exists19if not os.path.exists(ORIGINAL_DATASET):20    print(f"Error: The original dataset path '{ORIGINAL_DATASET}' does not exist.")21    exit(1)22 23try:24    # Create folders25    for split in ["train", "val", "test"]:26        for category in os.listdir(ORIGINAL_DATASET):27            os.makedirs(os.path.join(BASE_DIR, split, category), exist_ok=True)28 29    # Split data30    for category in os.listdir(ORIGINAL_DATASET):31        category_path = os.path.join(ORIGINAL_DATASET, category)32        images = os.listdir(category_path)33 34        train_imgs, temp_imgs = train_test_split(images, test_size=0.3, random_state=42)35        val_imgs, test_imgs = train_test_split(temp_imgs, test_size=0.33, random_state=42)  # 0.33 of 0.3 = 10%36 37        for img_name in train_imgs:38            shutil.copy(os.path.join(category_path, img_name), os.path.join(BASE_DIR, "train", category, img_name))39 40        for img_name in val_imgs:41            shutil.copy(os.path.join(category_path, img_name), os.path.join(BASE_DIR, "val", category, img_name))42 43        for img_name in test_imgs:44            shutil.copy(os.path.join(category_path, img_name), os.path.join(BASE_DIR, "test", category, img_name))45 46    print("Done splitting dataset!")47except Exception as e:48    print(f"An error occurred while splitting the dataset: {e}, please check the paths and try again.")49    exit(1)