35 lines
925 B
Python
35 lines
925 B
Python
# 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}")
|