import gradio as gr
from ultralytics import YOLO

# Load YOLOv8 model (downloads automatically on first run)
model = YOLO("yolov8n.pt")

def detect_objects(image, conf):
    if image is None:
        return None, "Please upload an image."

    # Run inference
    results = model(image, conf=conf)
    annotated_image = results[0].plot()

    # Count detected objects
    boxes = results[0].boxes
    if boxes is None or len(boxes) == 0:
        return annotated_image, "No objects detected."

    class_ids = boxes.cls.tolist()
    names = [model.names[int(i)] for i in class_ids]

    counts = {}
    for name in names:
        counts[name] = counts.get(name, 0) + 1

    summary = "Detected Objects:\n"
    for obj, count in counts.items():
        summary += f"{obj}: {count}\n"

    return annotated_image, summary


with gr.Blocks(title="YOLOv8 Image Recognition") as demo:
    gr.Markdown("## 🧠 YOLOv8 Image Recognition")
    gr.Markdown("Upload an image and detect objects automatically.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image")
            conf_slider = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.25,
                step=0.05,
                label="Confidence Threshold",
            )
            detect_button = gr.Button("🔍 Detect Objects")

        with gr.Column():
            image_output = gr.Image(label="Detected Image")
            text_output = gr.Textbox(label="Detection Summary")

    detect_button.click(
        fn=detect_objects,
        inputs=[image_input, conf_slider],
        outputs=[image_output, text_output],
    )

demo.launch()