jameslahm commited on
Commit
8b2cbe6
1 Parent(s): 6547fc1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +118 -0
app.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import PIL.Image as Image
2
+ import gradio as gr
3
+
4
+ from ultralytics import ASSETS, YOLOv10
5
+
6
+
7
+ def predict_image(img, model_id, image_size, conf_threshold):
8
+ model = YOLOv10.from_pretrained(f'jameslahm/{model_id}')
9
+ results = model.predict(
10
+ source=img,
11
+ conf=conf_threshold,
12
+ show_labels=True,
13
+ show_conf=True,
14
+ imgsz=image_size,
15
+ )
16
+
17
+ for r in results:
18
+ im_array = r.plot()
19
+ im = Image.fromarray(im_array[..., ::-1])
20
+
21
+ return im
22
+
23
+
24
+ def app():
25
+ with gr.Blocks():
26
+ with gr.Row():
27
+ with gr.Column():
28
+ image = gr.Image(type="pil", label="Image")
29
+
30
+ model_id = gr.Dropdown(
31
+ label="Model",
32
+ choices=[
33
+ "yolov10n",
34
+ "yolov10s",
35
+ "yolov10m",
36
+ "yolov10b",
37
+ "yolov10l",
38
+ "yolov10x",
39
+ ],
40
+ value="yolov10m",
41
+ )
42
+ image_size = gr.Slider(
43
+ label="Image Size",
44
+ minimum=320,
45
+ maximum=1280,
46
+ step=32,
47
+ value=640,
48
+ )
49
+ conf_threshold = gr.Slider(
50
+ label="Confidence Threshold",
51
+ minimum=0.1,
52
+ maximum=1.0,
53
+ step=0.1,
54
+ value=0.25,
55
+ )
56
+ yolov10_infer = gr.Button(value="Detect Objects")
57
+
58
+ with gr.Column():
59
+ output_image = gr.Image(type="pil", label="Annotated Image")
60
+
61
+ yolov10_infer.click(
62
+ fn=predict_image,
63
+ inputs=[
64
+ image,
65
+ model_id,
66
+ image_size,
67
+ conf_threshold,
68
+ ],
69
+ outputs=[output_image],
70
+ )
71
+
72
+ gr.Examples(
73
+ examples=[
74
+ [
75
+ "bug.jpg",
76
+ "yolov10x",
77
+ 640,
78
+ 0.25,
79
+ ],
80
+ [
81
+ "zidane.jpg",
82
+ "yolov10m",
83
+ 640,
84
+ 0.25,
85
+ ],
86
+ ],
87
+ fn=predict_image,
88
+ inputs=[
89
+ image,
90
+ model_id,
91
+ image_size,
92
+ conf_threshold,
93
+ iou_threshold,
94
+ ],
95
+ outputs=[output_image],
96
+ cache_examples=True,
97
+ )
98
+
99
+ gradio_app = gr.Blocks()
100
+ with gradio_app:
101
+ gr.HTML(
102
+ """
103
+ <h1 style='text-align: center'>
104
+ YOLOv10: Real-Time End-to-End Object Detection
105
+ </h1>
106
+ """)
107
+ gr.HTML(
108
+ """
109
+ <h3 style='text-align: center'>
110
+ YOLOv10: Real-Time End-to-End Object Detection
111
+ <a href='https://arxiv.org/abs/2405.14458' target='_blank'>arXiv</a> | <a href='https://github.com/THU-MIG/yolov10' target='_blank'>github</a>
112
+ </h3>
113
+ """)
114
+ with gr.Row():
115
+ with gr.Column():
116
+ app()
117
+ if __name__ == ''
118
+ gradio_app.launch(debug=True)