Spaces:
Sleeping
Sleeping
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) | |