import torch
import gradio as gr
from torchvision.transforms import functional as F
from PIL import Image

# Load the YOLOv8 model (assuming it is already converted to the Hugging Face format)
model = torch.hub.load('ultralytics/yolov8', 'custom', path='yolov5s.pt')

# Define the prediction function
def predict(image):
    # Preprocess the input image
    image_tensor = F.to_tensor(image)
    image_tensor.unsqueeze_(0)

    # Perform inference
    results = model(image_tensor)

    # Post-process the results
    # Extract the bounding box coordinates and class labels
    bboxes = results.xyxy[0].tolist()
    labels = results.names[0]

    return bboxes, labels

# Define the Gradio interface
inputs = gr.inputs.Image()
outputs = gr.outputs.Image()

interface = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, capture_session=True)

# Run the interface
interface.launch()