Spaces:
Runtime error
Runtime error
File size: 2,077 Bytes
0c65a85 | 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 | import os
import shutil
import random
from pathlib import Path
# Config
SOURCE_DIR = r"archive (1)\combined_images"
DATA_DIR = "data"
TRAIN_RATIO = 0.8
# Map source names to target names (if they differ slightly, or just lowercased)
# Source: MildDemented, ModerateDemented, NonDemented, VeryMildDemented
# Target: mild_demented, moderate_demented, nondemented, very_mild
CLASS_MAPPING = {
"MildDemented": "mild_demented",
"ModerateDemented": "moderate_demented",
"NonDemented": "nondemented",
"VeryMildDemented": "very_mild"
}
def setup_data():
if os.path.exists(DATA_DIR):
print(f"'{DATA_DIR}' already exists. Removing to ensure fresh setup...")
shutil.rmtree(DATA_DIR)
os.makedirs(os.path.join(DATA_DIR, "train"))
os.makedirs(os.path.join(DATA_DIR, "test"))
for source_name, target_name in CLASS_MAPPING.items():
src_class_path = os.path.join(SOURCE_DIR, source_name)
if not os.path.exists(src_class_path):
print(f"Warning: Source folder '{src_class_path}' not found.")
continue
# Create target class folders
os.makedirs(os.path.join(DATA_DIR, "train", target_name), exist_ok=True)
os.makedirs(os.path.join(DATA_DIR, "test", target_name), exist_ok=True)
images = [f for f in os.listdir(src_class_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
random.shuffle(images)
split_idx = int(len(images) * TRAIN_RATIO)
train_imgs = images[:split_idx]
test_imgs = images[split_idx:]
print(f"Processing {source_name}: {len(train_imgs)} train, {len(test_imgs)} test")
for img in train_imgs:
shutil.copy(os.path.join(src_class_path, img), os.path.join(DATA_DIR, "train", target_name, img))
for img in test_imgs:
shutil.copy(os.path.join(src_class_path, img), os.path.join(DATA_DIR, "test", target_name, img))
print("Data setup complete.")
if __name__ == "__main__":
setup_data()
|