Kushalmanda's picture
Update app.py
59ec4bf verified
raw
history blame
717 Bytes
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()