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

# load model
model = YOLO("best.pt")

def predict(image):
    results = model(image)
    r = results[0]
    
    # Group detections by class
    class_confidences = defaultdict(list)
    
    if len(r.boxes) > 0:
        for box in r.boxes:
            class_id = int(box.cls[0])
            class_name = model.names[class_id]
            confidence = float(box.conf[0]) * 100
            class_confidences[class_name].append(confidence)
    
    if class_confidences:
        lines = []
        sorted_classes = sorted(
            class_confidences.items(),
            key=lambda x: max(x[1]),
            reverse=True
        )
        
        for class_name, confidences in sorted_classes:
            avg_conf = round(sum(confidences) / len(confidences), 2)
            max_conf = round(max(confidences), 2)
            lines.append(f"• {class_name} — avg: {avg_conf}% | best: {max_conf}%")
        
        result_text = "\n".join(lines)
    else:
        result_text = "No detection"
    
    output_image = r.plot()
    return Image.fromarray(output_image), result_text

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Detection"),
        gr.Textbox(label="Result", lines=10)
    ],
    title="YOLOv8 Detection",
    description="Upload an image and detect objects"
)
demo.launch()