Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
|
3 |
+
import gradio as gr
|
4 |
+
import torch
|
5 |
+
from PIL import Image, ImageDraw, ImageFont
|
6 |
+
import numpy as np
|
7 |
+
import matplotlib.pyplot as plt
|
8 |
+
import os
|
9 |
+
|
10 |
+
# Download images
|
11 |
+
torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
|
12 |
+
torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', 'bus.jpg')
|
13 |
+
|
14 |
+
# Load YOLOv5 model
|
15 |
+
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
|
16 |
+
|
17 |
+
|
18 |
+
def yolo(im):
|
19 |
+
try:
|
20 |
+
# Check if the input is an Image object
|
21 |
+
if isinstance(im, Image.Image):
|
22 |
+
# Convert the PIL image to a numpy array
|
23 |
+
im_array = np.array(im)
|
24 |
+
|
25 |
+
# Perform inference with YOLOv5
|
26 |
+
results = model(im_array) # inference
|
27 |
+
|
28 |
+
# Get the bounding boxes and labels
|
29 |
+
boxes = results.xyxy[0].cpu().numpy()
|
30 |
+
|
31 |
+
# Convert the results to a PIL Image
|
32 |
+
output_image = Image.fromarray(im_array)
|
33 |
+
|
34 |
+
# Draw the bounding boxes and labels on the output image
|
35 |
+
draw = ImageDraw.Draw(output_image)
|
36 |
+
font = ImageFont.load_default(45)
|
37 |
+
|
38 |
+
for box in boxes:
|
39 |
+
label = results.names[int(box[5])]
|
40 |
+
draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline="red", width=3)
|
41 |
+
draw.text((box[0], box[1]), label, fill="blue", font=font)
|
42 |
+
|
43 |
+
return output_image
|
44 |
+
else:
|
45 |
+
raise ValueError("The input should be an Image object.")
|
46 |
+
except Exception as e:
|
47 |
+
print(f"Error processing image: {e}")
|
48 |
+
return None
|
49 |
+
|
50 |
+
|
51 |
+
# Define Gradio interface
|
52 |
+
inputs = gr.Image(type='pil', label="Original Image")
|
53 |
+
outputs = gr.Image(type="pil", label="Output Image")
|
54 |
+
title = "YOLOv5"
|
55 |
+
description = "YOLOv5 Gradio demo for object detection. Upload an image or click an example image to use."
|
56 |
+
|
57 |
+
article = "<p style='text-align: center'>YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset, and includes " \
|
58 |
+
"simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, " \
|
59 |
+
"and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> |" \
|
60 |
+
"<a href='https://apps.apple.com/app/id1452689527'>iOS App</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"
|
61 |
+
|
62 |
+
examples = [['zidane.jpg'], ['bus.jpg']]
|
63 |
+
|
64 |
+
gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch(debug=True)
|