# quick_labeler.py
import random
import shutil, os
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt

SOURCE = Path("alle_meine_fotos/")
images = list(SOURCE.glob("**/*.jpg")) + list(SOURCE.glob("**/*.png"))


DIRS = [
    "dataset/train/wallpaper", "dataset/train/no_wallpaper",
    "dataset/val/wallpaper",   "dataset/val/no_wallpaper",
]
for d in DIRS:
    Path(d).mkdir(parents=True, exist_ok=True)

for img_path in images:
    img = Image.open(img_path)
    plt.imshow(img)
    plt.title(img_path.name)
    plt.axis("off")
    plt.show(block=False)

    label = input("Wallpaper? (y/n/q): ").strip().lower()
    plt.close()

    if label == "q":
        break
    elif label in ("y", "n"):
        folder = "wallpaper" if label == "y" else "no_wallpaper"
        split  = "train" if random.random() < 0.8 else "val"
        shutil.copy(img_path, f"dataset/{split}/{folder}/{img_path.name}")