import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel
import torch.serialization

# Allowlist für YOLOv8-Modell
torch.serialization.add_safe_globals([DetectionModel])

# Lade YOLOv8-Modell
model = YOLO("yolov8n.pt")  # Nano-Modell, Standard-Laden

def analyze_image(image, prompt):
    # Prüfe Prompt
    if "what do you see" not in prompt.lower() and "was siehst du" not in prompt.lower():
        return "Nicht unterstützter Prompt. Verwende 'Was siehst du auf dem Bild?'"

    # Konvertiere PIL-Bild zu numpy-Format
    image_np = np.array(image)
    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    # YOLOv8 für Objekterkennung
    results = model.predict(source=image_np, conf=0.3, iou=0.5)
    objects = []
    
    for r in results:
        boxes = r.boxes.xyxy.cpu().numpy()
        labels = r.boxes.cls.cpu().numpy()
        class_names = model.names  # Klassen-Namen aus dem Modell
        for box, label in zip(boxes, labels):
            class_name = class_names[int(label)]
            objects.append(class_name)

    # Erstelle Textbeschreibung
    if not objects:
        return "Keine Objekte erkannt. Es sieht aus wie ein Chart oder ein anderes Bild ohne klar erkennbare Objekte."
    
    object_count = {}
    for obj in objects:
        object_count[obj] = object_count.get(obj, 0) + 1
    
    response = "Ich sehe folgende Objekte im Bild:\n"
    for obj, count in object_count.items():
        response += f"- {obj} ({count}x)\n"
    
    return response

# Erstelle Gradio-Schnittstelle
iface = gr.Interface(
    fn=analyze_image,
    inputs=[
        gr.Image(type="pil", label="Bild hochladen"),
        gr.Textbox(label="Prompt", placeholder="z. B. 'Was siehst du auf dem Bild?'")
    ],
    outputs="text",
    title="Einfache Bildanalyse mit YOLOv8",
    description="Lade ein Bild hoch und frage 'Was siehst du auf dem Bild?' für eine Beschreibung."
)

iface.launch()