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

# Load the YOLO model
model = YOLO('best.pt')

def predict(image):
    # Convert PIL Image to numpy array
    img_array = np.array(image)
    
    # Perform inference
    results = model(img_array)
    
    # Plot results
    for r in results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
    
    # Get detection counts and confidences
    class_counts = {i: 0 for i in range(len(model.names))}
    class_confidences = {i: [] for i in range(len(model.names))}
    
    for box in results[0].boxes:
        cls = int(box.cls[0])
        conf = float(box.conf[0])
        class_counts[cls] += 1
        class_confidences[cls].append(conf)
    
    # Create output string
    output = "Detection Results:\n"
    for i in range(len(model.names)):
        count = class_counts[i]
        avg_conf = np.mean(class_confidences[i]) if class_confidences[i] else 0
        output += f"{model.names[i]}: {count} detections (Avg. Confidence: {avg_conf:.2f})\n"
    
    return im, output

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Detected Image"),
        gr.Textbox(label="Detection Results")
    ],
    title="Helmet and License Plate Detection",
    description="Upload an image to detect helmets, license plates, and motorcyclists."
)

# Launch the interface
iface.launch()