Spaces:
Sleeping
Sleeping
File size: 2,706 Bytes
a549062 ca7b880 a549062 |
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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
# pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt gradio # install dependencies
import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import matplotlib.pyplot as plt
import os
# Download images
torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', 'bus.jpg')
# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
def yolo(im):
try:
# Check if the input is an Image object
if isinstance(im, Image.Image):
# Convert the PIL image to a numpy array
im_array = np.array(im)
# Perform inference with YOLOv5
results = model(im_array) # inference
# Get the bounding boxes and labels
boxes = results.xyxy[0].cpu().numpy()
# Convert the results to a PIL Image
output_image = Image.fromarray(im_array)
# Draw the bounding boxes and labels on the output image
draw = ImageDraw.Draw(output_image)
font = ImageFont.load_default(45)
for box in boxes:
label = results.names[int(box[5])]
draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline="red", width=3)
draw.text((box[0], box[1]), label, fill="blue", font=font)
return output_image
else:
raise ValueError("The input should be an Image object.")
except Exception as e:
print(f"Error processing image: {e}")
return None
# Define Gradio interface
inputs = gr.Image(type='pil', label="Original Image")
outputs = gr.Image(type="pil", label="Output Image")
title = "YOLOv5"
description = "YOLOv5 Gradio demo for object detection. Upload an image or click an example image to use."
article = "<p style='text-align: center'>YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset, and includes " \
"simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, " \
"and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> |" \
"<a href='https://apps.apple.com/app/id1452689527'>iOS App</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"
examples = [['zidane.jpg'], ['bus.jpg']]
gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch(debug=True)
|