import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO  # Import YOLO from ultralytics

# Load the YOLOv5 model
model = YOLO('yolov5s')  # Use the YOLOv5s pre-trained model

# Function to run inference on an image
def run_inference(image):
    # Convert the image from PIL format to a format compatible with OpenCV
    image = np.array(image)

    # Run YOLOv5 inference
    results = model.predict(source=image, save=False, conf=0.25, stream=False)

    # Annotate the image with detected objects
    annotated_image = results[0].plot()  # Use YOLO's built-in plotting function

    return annotated_image

# Create the Gradio interface
interface = gr.Interface(
    fn=run_inference,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="numpy"),
    title="YOLOv5 Object Detection",
    description="Upload an image to run YOLOv5 object detection and see the results."
)

# Launch the app
interface.launch()