import gradio as gr
import numpy as np
from ultralytics import YOLO
from PIL import Image
import cv2

# ----------------------------------
# Load Model (must be in Space root)
# ----------------------------------
model = YOLO("best.pt")

# Automatically use model's trained class names
CLASS_NAMES = model.names


# ----------------------------------
# Inference Function
# ----------------------------------
def detect_keys(image, conf_threshold):

    # Convert PIL → numpy
    img_np = np.array(image)

    # Run YOLO inference
    results = model(img_np, conf=float(conf_threshold))

    # Get annotated image with bounding boxes
    annotated = results[0].plot()

    # Convert BGR → RGB for Gradio display
    annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)

    # Extract detected class names
    detected_classes = set()
    if results[0].boxes is not None:
        for box in results[0].boxes:
            cls_id = int(box.cls[0])
            detected_classes.add(CLASS_NAMES[cls_id])

    if detected_classes:
        detected_text = "Detected: " + ", ".join(sorted(detected_classes))
    else:
        detected_text = "No keys detected"

    return annotated, detected_text


# ----------------------------------
# Gradio Interface
# ----------------------------------
demo = gr.Interface(
    fn=detect_keys,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(0.1, 1.0, value=0.25, step=0.05, label="Confidence Threshold"),
    ],
    outputs=[
        gr.Image(type="numpy", label="Detection Result"),
        gr.Textbox(label="Detected Labels"),
    ],
    title="🔑 Keys Detection (YOLO)",
    description="Upload an image to detect keys defects (cracking / missing).",
)

if __name__ == "__main__":
    demo.launch()