File size: 705 Bytes
0c3390c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
import torch
import gradio as gr
from PIL import Image
# Load pre-trained YOLOv5 model from ultralytics
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
def detect_objects(image):
# Perform inference
results = model(image)
results.render() # updates results.imgs with boxes and labels
detected_image = Image.fromarray(results.imgs[0])
return detected_image
# Define Gradio interface
iface = gr.Interface(
fn=detect_objects,
inputs=gr.inputs.Image(type="pil"),
outputs=gr.outputs.Image(type="pil"),
title="Object Detection with YOLOv5",
description="Upload an image and get the detected objects with their names."
)
# Launch the interface
iface.launch()
|