|
import os |
|
import shutil |
|
import random |
|
|
|
|
|
data_dir = "PlantVillage" |
|
train_dir = "dataset/train" |
|
test_dir = "dataset/test" |
|
split_ratio = 0.8 |
|
|
|
|
|
os.makedirs(train_dir, exist_ok=True) |
|
os.makedirs(test_dir, exist_ok=True) |
|
|
|
|
|
for category in os.listdir(data_dir): |
|
category_path = os.path.join(data_dir, category) |
|
|
|
if os.path.isdir(category_path): |
|
images = os.listdir(category_path) |
|
random.shuffle(images) |
|
|
|
split_index = int(len(images) * split_ratio) |
|
train_images = images[:split_index] |
|
test_images = images[split_index:] |
|
|
|
|
|
os.makedirs(os.path.join(train_dir, category), exist_ok=True) |
|
os.makedirs(os.path.join(test_dir, category), exist_ok=True) |
|
|
|
|
|
for img in train_images: |
|
shutil.move(os.path.join(category_path, img), os.path.join(train_dir, category, img)) |
|
for img in test_images: |
|
shutil.move(os.path.join(category_path, img), os.path.join(test_dir, category, img)) |
|
|
|
print("β
Dataset successfully split into train/test!") |
|
|