Koni commited on
Commit
c79c7e9
1 Parent(s): bf71e97

First commit trying to make this run

Browse files
Files changed (3) hide show
  1. .gitignore +1 -0
  2. app.py +200 -0
  3. requirements.txt +2 -2
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ venv
app.py ADDED
@@ -0,0 +1,200 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spaces
3
+ from huggingface_hub import hf_hub_download
4
+ import os
5
+
6
+ # make sure you have the following dependencies
7
+ import torch
8
+ import numpy as np
9
+ from models.common import DetectMultiBackend
10
+ from utils.general import non_max_suppression, scale_boxes
11
+ from utils.torch_utils import select_device, smart_inference_mode
12
+ from utils.augmentations import letterbox
13
+ import PIL.Image
14
+
15
+ #@smart_inference_mode()
16
+ @spaces.GPU
17
+ def yolov9_inference(img_path, model_id='YOLOv9-S_X_LOCO-converted.pt', img_size=640, conf_thres=0.1, iou_thres=0.4):
18
+ """
19
+ Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust
20
+ the input size and apply test time augmentation.
21
+
22
+ :param model_path: Path to the YOLOv9 model file.
23
+ :param conf_threshold: Confidence threshold for NMS.
24
+ :param iou_threshold: IoU threshold for NMS.
25
+ :param img_path: Path to the image file.
26
+ :param size: Optional, input size for inference.
27
+ :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
28
+ """
29
+
30
+ # Load the model
31
+ model_path = download_models(model_id)
32
+
33
+ # Initialize
34
+ device = select_device('0')
35
+ model = DetectMultiBackend(model_path, device="0", fp16=False, data='data/coco.yaml')
36
+ stride, names, pt = model.stride, model.names, model.pt
37
+
38
+ # Load image
39
+ img = np.array(PIL.Image.open(img_path))
40
+ img = letterbox(img0, img_size, stride=stride, auto=True)[0]
41
+ img = img[:, :, ::-1].transpose(2, 0, 1)
42
+ img = np.ascontiguousarray(img)
43
+ img = torch.from_numpy(img).to(device).float()
44
+ img /= 255.0
45
+ if img.ndimension() == 3:
46
+ img = img.unsqueeze(0)
47
+
48
+ # Inference
49
+ results = model(img, augment=False, visualize=False)
50
+
51
+ # Apply NMS
52
+ results = non_max_suppression(results[0][0], conf_thres, iou_thres, classes=None, max_det=1000)
53
+
54
+ output = results.render()
55
+
56
+ return output[0]
57
+
58
+
59
+
60
+ def download_models(model_id):
61
+ hf_hub_download("KoniHD/LOCO-Detection", filename=f"{model_id}", local_dir=f"./",
62
+ token=os.getenv("HF_TOKEN"))
63
+ return f"./{model_id}"
64
+
65
+ # @spaces.GPU
66
+ # def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_threshold):
67
+ # """
68
+ # Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust
69
+ # the input size and apply test time augmentation.
70
+
71
+ # :param model_path: Path to the YOLOv9 model file.
72
+ # :param conf_threshold: Confidence threshold for NMS.
73
+ # :param iou_threshold: IoU threshold for NMS.
74
+ # :param img_path: Path to the image file.
75
+ # :param size: Optional, input size for inference.
76
+ # :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
77
+ # """
78
+ # # Import YOLOv9
79
+ # import yolov9
80
+
81
+ # # Load the model
82
+ # model_path = download_models(model_id)
83
+ # model = yolov9.load(model_path, device="cuda:0")
84
+
85
+ # # Set model parameters
86
+ # model.conf = conf_threshold
87
+ # model.iou = iou_threshold
88
+
89
+ # # Perform inference
90
+ # results = model(img_path, size=image_size)
91
+
92
+ # # Optionally, show detection bounding boxes on image
93
+ # output = results.render()
94
+
95
+ # return output[0]
96
+
97
+
98
+ def app():
99
+ with gr.Blocks():
100
+ with gr.Row():
101
+ with gr.Column():
102
+ img_path = gr.Image(type="filepath", label="Image")
103
+ model_path = gr.Dropdown(
104
+ label="Model",
105
+ choices=[
106
+ "YOLOv9-S_X_LOCO-converted.pt",
107
+ "YOLOv9_S_X_LOCO.pt",
108
+ "YOLOv9-E_X_LOCO-converted.pt",
109
+ "YOLOv9_E_X_LOCO.pt",
110
+ ],
111
+ value="YOLOv9-S_X_LOCO-converted.pt",
112
+ )
113
+ image_size = gr.Slider(
114
+ label="Image Size",
115
+ minimum=320,
116
+ maximum=1280,
117
+ step=32,
118
+ value=640,
119
+ )
120
+ conf_threshold = gr.Slider(
121
+ label="Confidence Threshold",
122
+ minimum=0.1,
123
+ maximum=1.0,
124
+ step=0.1,
125
+ value=0.4,
126
+ )
127
+ iou_threshold = gr.Slider(
128
+ label="IoU Threshold",
129
+ minimum=0.1,
130
+ maximum=1.0,
131
+ step=0.1,
132
+ value=0.5,
133
+ )
134
+ yolov9_infer = gr.Button(value="Inference")
135
+
136
+ with gr.Column():
137
+ output_numpy = gr.Image(type="numpy",label="Output")
138
+
139
+ yolov9_infer.click(
140
+ fn=yolov9_inference,
141
+ inputs=[
142
+ img_path,
143
+ model_path,
144
+ image_size,
145
+ conf_threshold,
146
+ iou_threshold,
147
+ ],
148
+ outputs=[output_numpy],
149
+ )
150
+
151
+ gr.Examples(
152
+ examples=[
153
+ [
154
+ "data/zidane.jpg",
155
+ "YOLOv9-S_X_LOCO-converted.pt",
156
+ 640,
157
+ 0.4,
158
+ 0.5,
159
+ ],
160
+ [
161
+ "data/huggingface.jpg",
162
+ "YOLOv9-E_X_LOCO-converted.pt",
163
+ 640,
164
+ 0.4,
165
+ 0.5,
166
+ ],
167
+ ],
168
+ fn=yolov9_inference,
169
+ inputs=[
170
+ img_path,
171
+ model_path,
172
+ image_size,
173
+ conf_threshold,
174
+ iou_threshold,
175
+ ],
176
+ outputs=[output_numpy],
177
+ cache_examples=True,
178
+ )
179
+
180
+
181
+ gradio_app = gr.Blocks()
182
+ with gradio_app:
183
+ gr.HTML(
184
+ """
185
+ <h1 style='text-align: center'>
186
+ YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information
187
+ </h1>
188
+ """)
189
+ gr.HTML(
190
+ """
191
+ <h3 style='text-align: center'>
192
+ Follow me for more!
193
+ <a href='https://twitter.com/konihd_7' target='_blank'>Twitter</a> | <a href='https://github.com/KoniHD' target='_blank'>Github</a> | <a href='https://www.linkedin.com/in/konstantin-zeck/' target='_blank'>Linkedin</a> | <a href='https://www.huggingface.co/KoniHD/' target='_blank'>HuggingFace</a>
194
+ </h3>
195
+ """)
196
+ with gr.Row():
197
+ with gr.Column():
198
+ app()
199
+
200
+ gradio_app.launch(debug=True)
requirements.txt CHANGED
@@ -1,8 +1,6 @@
1
  # requirements
2
  # Usage: pip install -r requirements.txt
3
 
4
- git+https://github.com/WongKinYiu/yolov9.git@main
5
-
6
  # Base ------------------------------------------------------------------------
7
  gitpython
8
  ipython
@@ -42,6 +40,8 @@ seaborn>=0.11.0
42
 
43
  # Deploy ----------------------------------------------------------------------
44
  # tritonclient[all]~=2.24.0
 
 
45
 
46
  # Extras ----------------------------------------------------------------------
47
  # mss
 
1
  # requirements
2
  # Usage: pip install -r requirements.txt
3
 
 
 
4
  # Base ------------------------------------------------------------------------
5
  gitpython
6
  ipython
 
40
 
41
  # Deploy ----------------------------------------------------------------------
42
  # tritonclient[all]~=2.24.0
43
+ git+https://github.com/KoniHD/yolov9.git@main#egg=yolov9
44
+ huggingface_hub
45
 
46
  # Extras ----------------------------------------------------------------------
47
  # mss