import torch
import gradio as gr
from PIL import Image

# Load the YOLOv5 model from the uploaded file (e.g., 'yolov5s.pt')
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt')  # Adjust the file name if needed

# Define the function for image inference
def predict_image(image):
    results = model(image)  # Run inference on the input image
    results.show()  # Optionally visualize the results (bounding boxes)
    return results.pandas().xywh  # Return the results (bounding box coordinates)

# Set up Gradio interface to allow image uploads and get predictions
interface = gr.Interface(fn=predict_image, inputs=gr.Image(), outputs=gr.Dataframe())

# Launch the interface
interface.launch()