Spaces:
Runtime error
Runtime error
File size: 1,184 Bytes
3e1cfd6 2f46953 3e1cfd6 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
# Load the YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
# Function to perform object detection
def detect_objects(image):
# Convert the image to a numpy array
img_array = np.array(image)
# Perform object detection
results = model(img_array)
# Draw bounding boxes and labels
for index, row in results.pandas().xyxy[0].iterrows():
x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
label = row['name']
cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(img_array, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
# Convert the numpy array back to an image
output_image = Image.fromarray(img_array)
return output_image
# Set up the Gradio interface
interface = gr.Interface(
fn=detect_objects,
inputs=gr.inputs.Image(type="pil", tool="editor"),
outputs=gr.outputs.Image(type="pil"),
live=True,
title="Real-time Object Detection with YOLOv5"
)
# Launch the interface
interface.launch()
|