# Import required libraries
import torch
import gradio as gr
from PIL import Image

# Load the pretrained YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True)

# Function to process the image and return detections
def detect_objects(image):
    # Perform inference on the uploaded image
    results = model(image)
    
    # Plot results on the image (YOLOv5 provides results with bounding boxes, class names, and confidence scores)
    results_img = results.render()[0]  # Render the detections on the image
    
    # Convert to a PIL Image for compatibility with Gradio
    return Image.fromarray(results_img)

# Define the Gradio interface
interface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Object Detection App",
    description="Upload an image to detect objects using the YOLOv5 model."
)

# Launch the Gradio app
interface.launch()