Spaces:
Build error
Build error
sagemaker
commited on
Commit
•
7e3d7f8
1
Parent(s):
9ee3e39
init
Browse files
README.md
CHANGED
@@ -1,6 +1,6 @@
|
|
1 |
---
|
2 |
title: YOLOS Object Detection
|
3 |
-
emoji:
|
4 |
colorFrom: pink
|
5 |
colorTo: purple
|
6 |
sdk: gradio
|
|
|
1 |
---
|
2 |
title: YOLOS Object Detection
|
3 |
+
emoji: 👤
|
4 |
colorFrom: pink
|
5 |
colorTo: purple
|
6 |
sdk: gradio
|
app.py
ADDED
@@ -0,0 +1,70 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import AutoFeatureExtractor, YolosForObjectDetection
|
2 |
+
import gradio as gr
|
3 |
+
from PIL import Image
|
4 |
+
import torch
|
5 |
+
import matplotlib.pyplot as plt
|
6 |
+
import io
|
7 |
+
|
8 |
+
|
9 |
+
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
|
10 |
+
[0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
|
11 |
+
|
12 |
+
|
13 |
+
def infer(img, model_name):
|
14 |
+
feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/{model_name}")
|
15 |
+
model = YolosForObjectDetection.from_pretrained(f"hustvl/{model_name}")
|
16 |
+
|
17 |
+
img = Image.fromarray(img)
|
18 |
+
|
19 |
+
pixel_values = feature_extractor(img, return_tensors="pt").pixel_values
|
20 |
+
|
21 |
+
with torch.no_grad():
|
22 |
+
outputs = model(pixel_values, output_attentions=True)
|
23 |
+
|
24 |
+
probas = outputs.logits.softmax(-1)[0, :, :-1]
|
25 |
+
keep = probas.max(-1).values > 0.9
|
26 |
+
|
27 |
+
target_sizes = torch.tensor(img.size[::-1]).unsqueeze(0)
|
28 |
+
postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
|
29 |
+
bboxes_scaled = postprocessed_outputs[0]['boxes']
|
30 |
+
|
31 |
+
res_img = plot_results(img, probas[keep], bboxes_scaled[keep], model)
|
32 |
+
|
33 |
+
return res_img
|
34 |
+
|
35 |
+
def plot_results(pil_img, prob, boxes, model):
|
36 |
+
plt.figure(figsize=(16,10))
|
37 |
+
plt.imshow(pil_img)
|
38 |
+
ax = plt.gca()
|
39 |
+
colors = COLORS * 100
|
40 |
+
for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
|
41 |
+
ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
|
42 |
+
fill=False, color=c, linewidth=3))
|
43 |
+
cl = p.argmax()
|
44 |
+
text = f'{model.config.id2label[cl.item()]}: {p[cl]:0.2f}'
|
45 |
+
ax.text(xmin, ymin, text, fontsize=15,
|
46 |
+
bbox=dict(facecolor='yellow', alpha=0.5))
|
47 |
+
plt.axis('off')
|
48 |
+
return fig2img(plt.gcf())
|
49 |
+
|
50 |
+
def fig2img(fig):
|
51 |
+
buf = io.BytesIO()
|
52 |
+
fig.savefig(buf)
|
53 |
+
buf.seek(0)
|
54 |
+
img = Image.open(buf)
|
55 |
+
return img
|
56 |
+
|
57 |
+
description = """Object Detection with YOLOS. Choose your model and you're good to go."""
|
58 |
+
|
59 |
+
image_in = gr.components.Image()
|
60 |
+
image_out = gr.components.Image()
|
61 |
+
model_choice = gr.components.Dropdown(["yolos-tiny", "yolos-small", "yolos_base", "yolos-small-300", "yolos-small-dwr"], value="yolos-small")
|
62 |
+
|
63 |
+
Iface = gr.Interface(
|
64 |
+
fn=infer,
|
65 |
+
inputs=[image_in,model_choice],
|
66 |
+
outputs=image_out,
|
67 |
+
examples=[["examples/10_People_Marching_People_Marching_2_120.jpg"], ["examples/12_Group_Group_12_Group_Group_12_26.jpg"], ["examples/43_Row_Boat_Canoe_43_247.jpg"]],
|
68 |
+
title="Object Detection with YOLOS",
|
69 |
+
description=description,
|
70 |
+
).launch()
|
examples/10_People_Marching_People_Marching_2_120.jpg
ADDED
examples/12_Group_Group_12_Group_Group_12_26.jpg
ADDED
examples/43_Row_Boat_Canoe_43_247.jpg
ADDED
requirements.txt
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
transformers
|
2 |
+
pillow
|
3 |
+
torch
|
4 |
+
matplotlib
|