import gradio as gr
from ultralytics import YOLO
from PIL import Image

# -----------------------------
# Load YOLO model
# -----------------------------
model = YOLO("./data/best.pt")  # make sure this path matches your folder structure

# -----------------------------
# Prediction function
# -----------------------------
def predict(image):
    # Run prediction
    results = model.predict(image, conf=0.5)

    # Annotated image with bounding boxes
    result_img = results[0].plot()

    # Extract detected labels
    detected_labels = results[0].boxes.cls.tolist()
    names = results[0].names
    detected_objects = [names[int(cls_id)] for cls_id in detected_labels]

    # Text output
    if detected_objects:
        label_text = f"✅ Detected objects: {', '.join(detected_objects)}"
    else:
        label_text = "❌ No objects detected."

    return Image.fromarray(result_img), label_text

# -----------------------------
# Gradio Interface
# -----------------------------
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil", label="Detection Result"), gr.Textbox(label="Detected Objects")],
    title="🥤 Bottle Detection with YOLOv11",
    description="Upload an image to check if a **bottle** is detected using your trained YOLOv11 model."
)

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