musfiqdehan commited on
Commit
1c6497f
1 Parent(s): 67420f1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +173 -0
app.py ADDED
@@ -0,0 +1,173 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spaces
3
+ from huggingface_hub import hf_hub_download
4
+
5
+
6
+ def download_models(model_id):
7
+ hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir=f"./")
8
+ return f"./{model_id}"
9
+
10
+ @spaces.GPU
11
+ def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_threshold):
12
+ """
13
+ Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust
14
+ the input size and apply test time augmentation.
15
+
16
+ :param model_path: Path to the YOLOv9 model file.
17
+ :param conf_threshold: Confidence threshold for NMS.
18
+ :param iou_threshold: IoU threshold for NMS.
19
+ :param img_path: Path to the image file.
20
+ :param size: Optional, input size for inference.
21
+ :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
22
+ """
23
+ # Import YOLOv9
24
+ import yolov9
25
+
26
+ # Load the model
27
+ model_path = download_models(model_id)
28
+ model = yolov9.load(model_path, device="cuda:0")
29
+
30
+ # Set model parameters
31
+ model.conf = conf_threshold
32
+ model.iou = iou_threshold
33
+
34
+ # Perform inference
35
+ results = model(img_path, size=image_size)
36
+
37
+ # Optionally, show detection bounding boxes on image
38
+ output = results.render()
39
+
40
+ return output[0]
41
+
42
+
43
+ def app():
44
+ with gr.Blocks():
45
+ with gr.Row():
46
+ with gr.Column():
47
+ img_path = gr.Image(type="filepath", label="Image")
48
+ model_path = gr.Dropdown(
49
+ label="Model",
50
+ choices=[
51
+ "gelan-c.pt",
52
+ "gelan-e.pt",
53
+ "yolov9-c.pt",
54
+ "yolov9-e.pt",
55
+ ],
56
+ value="gelan-e.pt",
57
+ )
58
+ image_size = gr.Slider(
59
+ label="Image Size",
60
+ minimum=320,
61
+ maximum=1280,
62
+ step=32,
63
+ value=640,
64
+ )
65
+ conf_threshold = gr.Slider(
66
+ label="Confidence Threshold",
67
+ minimum=0.1,
68
+ maximum=1.0,
69
+ step=0.1,
70
+ value=0.4,
71
+ )
72
+ iou_threshold = gr.Slider(
73
+ label="IoU Threshold",
74
+ minimum=0.1,
75
+ maximum=1.0,
76
+ step=0.1,
77
+ value=0.5,
78
+ )
79
+ yolov9_infer = gr.Button(value="Inference")
80
+
81
+ with gr.Column():
82
+ output_numpy = gr.Image(type="numpy",label="Output")
83
+
84
+ yolov9_infer.click(
85
+ fn=yolov9_inference,
86
+ inputs=[
87
+ img_path,
88
+ model_path,
89
+ image_size,
90
+ conf_threshold,
91
+ iou_threshold,
92
+ ],
93
+ outputs=[output_numpy],
94
+ )
95
+
96
+ gr.Examples(
97
+ examples=[
98
+ [
99
+ "example-data/img-1.jpg",
100
+ "gelan-e.pt",
101
+ 640,
102
+ 0.4,
103
+ 0.5,
104
+ ],
105
+ [
106
+ "example-data/img-2.jpg",
107
+ "yolov9-c.pt",
108
+ 640,
109
+ 0.4,
110
+ 0.5,
111
+ ],
112
+ [
113
+ "example-data/img-3.jpg",
114
+ "yolov9-c.pt",
115
+ 640,
116
+ 0.4,
117
+ 0.5,
118
+ ],
119
+ [
120
+ "example-data/img-4.jpg",
121
+ "yolov9-e.pt",
122
+ 640,
123
+ 0.4,
124
+ 0.5,
125
+ ],
126
+ [
127
+ "example-data/img-5.jpg",
128
+ "gelan-e.pt",
129
+ 740,
130
+ 0.4,
131
+ 0.5,
132
+ ],
133
+ [
134
+ "example-data/img-6.jpg",
135
+ "yolov9-c.pt",
136
+ 640,
137
+ 0.4,
138
+ 0.5,
139
+ ],
140
+ [
141
+ "example-data/img-4.jpg",
142
+ "gelan-c.pt",
143
+ 640,
144
+ 0.4,
145
+ 0.5,
146
+ ],
147
+ ],
148
+ fn=yolov9_inference,
149
+ inputs=[
150
+ img_path,
151
+ model_path,
152
+ image_size,
153
+ conf_threshold,
154
+ iou_threshold,
155
+ ],
156
+ outputs=[output_numpy],
157
+ cache_examples=True,
158
+ )
159
+
160
+
161
+ gradio_app = gr.Blocks()
162
+ with gradio_app:
163
+ gr.HTML(
164
+ """
165
+ <h1 style='text-align: center'>
166
+ Object Detection Using YOLO
167
+ </h1>
168
+ """)
169
+ with gr.Row():
170
+ with gr.Column():
171
+ app()
172
+
173
+ gradio_app.launch(debug=True)