khoaliamle/Garbage_Classification_YOLO

1import os2import shutil3import random4 5# Define paths6train_dir = "train"7val_dir = "val"8test_dir = "test"9 10# Define split ratios11train_ratio = 0.812val_ratio = 0.113test_ratio = 0.114 15# Ensure output directories exist16for split_dir in [train_dir, val_dir, test_dir]:17    os.makedirs(split_dir, exist_ok=True)18 19# Get class names (subdirectories current dir)20class_names = [d for d in os.listdir() if os.path.isdir(d) and d not in {"train", "val", "test"}]21 22# Process each class23for class_name in class_names:24    class_path = class_name25    images = [f for f in os.listdir(class_path) if os.path.isfile(os.path.join(class_path, f))]26    27    # Shuffle images randomly28    random.shuffle(images)29 30    # Compute split indices31    total_images = len(images)32    train_count = int(total_images * train_ratio)33    val_count = int(total_images * val_ratio)34 35    # Split images36    train_images = images[:train_count]37    val_images = images[train_count:train_count + val_count]38    test_images = images[train_count + val_count:]39 40    # Define destination directories for the class41    for split_name, split_images in zip(["train", "val", "test"], [train_images, val_images, test_images]):42        split_class_dir = os.path.join(split_name, class_name)43        os.makedirs(split_class_dir, exist_ok=True)44        45        # Move images46        for image in split_images:47            src = os.path.join(class_path, image)48            dst = os.path.join(split_class_dir, image)49            shutil.move(src, dst)50 51print("Dataset successfully split into train, val, and test sets.")52