qbcode commited on
Commit
04e2f0f
·
verified ·
1 Parent(s): e6efc44

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +147 -0
app.py ADDED
@@ -0,0 +1,147 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spaces
3
+ from huggingface_hub import hf_hub_download
4
+
5
+ def download_models(model_id):
6
+ hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir=f"./")
7
+ return f"./{model_id}"
8
+
9
+ @spaces.GPU
10
+ def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_threshold):
11
+ """
12
+ Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust
13
+ the input size and apply test time augmentation.
14
+
15
+ :param model_path: Path to the YOLOv9 model file.
16
+ :param conf_threshold: Confidence threshold for NMS.
17
+ :param iou_threshold: IoU threshold for NMS.
18
+ :param img_path: Path to the image file.
19
+ :param size: Optional, input size for inference.
20
+ :return: Annotated image with the number of objects detected.
21
+ """
22
+ # Import YOLOv9
23
+ import yolov9
24
+
25
+ # Load the model
26
+ model_path = download_models(model_id)
27
+ model = yolov9.load(model_path, device="cuda:0")
28
+
29
+ # Set model parameters
30
+ model.conf = conf_threshold
31
+ model.iou = iou_threshold
32
+
33
+ # Perform inference
34
+ results = model(img_path, size=image_size)
35
+
36
+ # Optionally, show detection bounding boxes on image
37
+ annotated_image = results.render()
38
+
39
+ # Count the number of detected objects
40
+ num_objects_detected = len(results.pred[0])
41
+
42
+ return annotated_image, f"Number of objects detected: {num_objects_detected}"
43
+
44
+ def app():
45
+ with gr.Blocks():
46
+ with gr.Row():
47
+ with gr.Column():
48
+ img_path = gr.Image(type="filepath", label="Image")
49
+ model_path = gr.Dropdown(
50
+ label="Model",
51
+ choices=[
52
+ "gelan-c.pt",
53
+ "gelan-e.pt",
54
+ "yolov9-c.pt",
55
+ "yolov9-e.pt",
56
+ ],
57
+ value="gelan-e.pt",
58
+ )
59
+ image_size = gr.Slider(
60
+ label="Image Size",
61
+ minimum=320,
62
+ maximum=1280,
63
+ step=32,
64
+ value=640,
65
+ )
66
+ conf_threshold = gr.Slider(
67
+ label="Confidence Threshold",
68
+ minimum=0.1,
69
+ maximum=1.0,
70
+ step=0.1,
71
+ value=0.4,
72
+ )
73
+ iou_threshold = gr.Slider(
74
+ label="IoU Threshold",
75
+ minimum=0.1,
76
+ maximum=1.0,
77
+ step=0.1,
78
+ value=0.5,
79
+ )
80
+ yolov9_infer = gr.Button(value="Inference")
81
+
82
+ with gr.Column():
83
+ output_image = gr.Image(type="numpy",label="Output")
84
+ num_objects_detected = gr.Textbox(label="Number of Objects Detected", readonly=True)
85
+
86
+ yolov9_infer.click(
87
+ fn=yolov9_inference,
88
+ inputs=[
89
+ img_path,
90
+ model_path,
91
+ image_size,
92
+ conf_threshold,
93
+ iou_threshold,
94
+ ],
95
+ outputs=[output_image, num_objects_detected],
96
+ )
97
+
98
+ gr.Examples(
99
+ examples=[
100
+ [
101
+ "data/zidane.jpg",
102
+ "gelan-e.pt",
103
+ 640,
104
+ 0.4,
105
+ 0.5,
106
+ ],
107
+ [
108
+ "data/huggingface.jpg",
109
+ "yolov9-c.pt",
110
+ 640,
111
+ 0.4,
112
+ 0.5,
113
+ ],
114
+ ],
115
+ fn=yolov9_inference,
116
+ inputs=[
117
+ img_path,
118
+ model_path,
119
+ image_size,
120
+ conf_threshold,
121
+ iou_threshold,
122
+ ],
123
+ outputs=[output_image, num_objects_detected],
124
+ cache_examples=True,
125
+ )
126
+
127
+
128
+ gradio_app = gr.Blocks()
129
+ with gradio_app:
130
+ gr.HTML(
131
+ """
132
+ <h1 style='text-align: center'>
133
+ YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information
134
+ </h1>
135
+ """)
136
+ gr.HTML(
137
+ """
138
+ <h3 style='text-align: center'>
139
+ Follow me for more!
140
+ <a href='https://twitter.com/kadirnar_ai' target='_blank'>Twitter</a> | <a href='https://github.com/kadirnar' target='_blank'>Github</a> | <a href='https://www.linkedin.com/in/kadir-nar/' target='_blank'>Linkedin</a> | <a href='https://www.huggingface.co/kadirnar/' target='_blank'>HuggingFace</a>
141
+ </h3>
142
+ """)
143
+ with gr.Row():
144
+ with gr.Column():
145
+ app()
146
+
147
+ gradio_app.launch(debug=True)