sagemaker commited on
Commit
7e3d7f8
1 Parent(s): 9ee3e39
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