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

# Load YOLO model
model = YOLO("yolov8n.pt")  # Replace with your own model if needed

def detect_objects(img):
    results = model(img)
    annotated_frame = results[0].plot()

    detections = []
    for box in results[0].boxes:
        cls_id = int(box.cls[0].item())
        label = model.names[cls_id]
        conf = float(box.conf[0].item())
        detections.append((label, conf))

    # Sort detections by confidence
    detections.sort(key=lambda x: x[1], reverse=True)

    # Header message if top detection is "hot dog"
    if detections and detections[0][0].lower() == "hot dog":
        detection_text = "🌭 this is hot dog!"
    else:
        detection_text = "this is not hot dog!"

    # Add main detection summary
    #detection_text += "Top Detections:\n"
    #detection_text += "\n".join([f"{label}: {conf:.2f}" for label, conf in detections[:5]])

    # Add full list of detections below
    detection_text += "\n\nDetected objects:\n"
    for label, conf in detections:
        detection_text += f"- {label}: {conf:.2f}\n"

    return Image.fromarray(annotated_frame), detection_text

# Gradio interface
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Detected Image"),
        gr.Textbox(label="Detected Objects")
    ],
    title="YOLO Object Detection",
    description="Upload an image to detect objects using YOLO. If the top result is a hot dog, it will tell you 🌭"
)

iface.launch(share=True)