# INSTALLATION PYTHON - ANFANG #
# 1. https://www.python.org/downloads/windows/
# 2. pip install opencv-python torch torchvision matplotlib ultralytics
# INSTALLATION PYTHON - ENDE #

# EINSTELLUNGEN - ANFANG #
# Suchbegriffe
filter_keywords = ['car', 'street', 'tree']

# Fotordner bzw. Pfad
image_dir = r'images'

# Ordnername für gefundene Treffer
treffer_dir = os.path.join(image_dir, 'Treffer')
# EINSTELLUNGEN - ENDE #

# START SCRIPT #
import os
import shutil
from ultralytics import YOLO

# Ordner erstellen, falls nicht vorhanden
os.makedirs(treffer_dir, exist_ok=True)

# Lade YOLOv8n-Modell
model = YOLO('yolov8n.pt')

# Durchlaufe Bilder
for filename in os.listdir(image_dir):
	if not filename.lower().endswith(('.jpg', '.jpeg', '.png')):
		continue

	image_path = os.path.join(image_dir, filename)
	results = model(image_path)

	# Alle erkannten Labels
	labels = [model.model.names[int(cls)] for cls in results[0].boxes.cls]
	print(f'{filename}: erkannte Labels: {labels}')

	# Falls ein Filterbegriff vorkommt → verschieben
	if any(keyword in labels for keyword in filter_keywords):
		zielpfad = os.path.join(treffer_dir, filename)
		print(f'✔ {filename} → {zielpfad} ({labels})')
		shutil.move(image_path, zielpfad)
	else:
		print(f'✘ {filename} enthält keine Treffer.')