split.py

import os
import random
import shutil

# ---------------- CONFIG ---------------- #
SOURCE_DIR = "mobilenet_data"
TARGET_DIR = "dataset"

TRAIN_RATIO = 0.8      # 80% train, 20% val
RANDOM_SEED = 42       # reproducibility
# ---------------------------------------- #


def split_train_val():
    random.seed(RANDOM_SEED)

    classes = [
        d for d in os.listdir(SOURCE_DIR)
        if os.path.isdir(os.path.join(SOURCE_DIR, d))
    ]

    if not classes:
        raise RuntimeError("No class folders found. Check SOURCE_DIR.")

    for split in ["train", "val"]:
        for cls in classes:
            os.makedirs(os.path.join(TARGET_DIR, split, cls), exist_ok=True)

    for cls in classes:
        cls_path = os.path.join(SOURCE_DIR, cls)
        images = [
            f for f in os.listdir(cls_path)
            if f.lower().endswith((".jpg", ".jpeg", ".png"))
        ]

        if len(images) < 5:
            print(f"⚠️  Very few samples in class '{cls}'")

        random.shuffle(images)

        split_idx = int(len(images) * TRAIN_RATIO)
        train_imgs = images[:split_idx]
        val_imgs = images[split_idx:]

        for img in train_imgs:
            shutil.copy2(
                os.path.join(cls_path, img),
                os.path.join(TARGET_DIR, "train", cls, img)
            )

        for img in val_imgs:
            shutil.copy2(
                os.path.join(cls_path, img),
                os.path.join(TARGET_DIR, "val", cls, img)
            )

        print(
            f"Class '{cls}': "
            f"{len(train_imgs)} train | {len(val_imgs)} val"
        )

    print("\n✅ Train/Val split completed successfully.")


if __name__ == "__main__":
    split_train_val()