Spaces:
Sleeping
Sleeping
File size: 2,459 Bytes
1d971c4 367ee7b 1d971c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import os
import shutil
import random
random.seed(42)
def prepare_data(source_dir, target_dir, classes_to_include):
"""
Restructures the dataset into binary classes (fresh/spoiled).
"""
for label in ['fresh', 'spoiled']:
os.makedirs(os.path.join(target_dir, label), exist_ok=True)
for category in os.listdir(source_dir):
# Determine if this category belongs to 'fresh' or 'spoiled'
if any(incl in category for incl in classes_to_include):
label = 'fresh' if 'fresh' in category else 'spoiled'
source_path = os.path.join(source_dir, category)
for img in os.listdir(source_path):
# Ensure we only copy image files (basic check)
if img.lower().endswith(('.png', '.jpg', '.jpeg')):
shutil.copy(
os.path.join(source_path, img),
os.path.join(target_dir, label, f"{category}_{img}")
)
def split_data(processed_dir, split_dir, train_p=0.7, val_p=0.15):
"""
Splits the processed binary data into train, val, and test sets.
"""
for split in ['train', 'val', 'test']:
for label in ['fresh', 'spoiled']:
os.makedirs(os.path.join(split_dir, split, label), exist_ok=True)
for label in ['fresh', 'spoiled']:
label_dir = os.path.join(processed_dir, label)
images = os.listdir(label_dir)
random.shuffle(images)
n = len(images)
train_end = int(n * train_p)
val_end = int(n * (train_p + val_p))
splits = {
'train': images[:train_end],
'val': images[train_end:val_end],
'test': images[val_end:]
}
for split_name, split_images in splits.items():
for img in split_images:
shutil.copy(
os.path.join(label_dir, img),
os.path.join(split_dir, split_name, label, img)
)
if __name__ == "__main__":
SOURCE = "./data/raw/dataset"
PROCESSED = "./data/processed"
FINAL = "./data/split"
# Strictly fruits and vegetables as requested "for now"
CLASSES = ['fruits', 'vegetables']
print(f"Restructuring data for: {CLASSES}...")
prepare_data(SOURCE, PROCESSED, CLASSES)
print("Splitting data into train/val/test...")
split_data(PROCESSED, FINAL)
print("Data preparation complete.")
|