kadirnar commited on
Commit
80dd37b
·
verified ·
1 Parent(s): 60c1a7b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -8
app.py CHANGED
@@ -4,9 +4,23 @@ import os
4
  from huggingface_hub import hf_hub_download
5
 
6
 
7
- def download_models(model_id):
8
- hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir=f"./{model_id}")
9
- return f"./{model_id}"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
 
11
 
12
  @spaces.GPU
@@ -26,7 +40,7 @@ def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_thresho
26
  import yolov9
27
 
28
  # Load the model
29
- model_path = download_models(model_id)
30
  model = yolov9.load(model_path, device="cuda")
31
 
32
  # Set model parameters
@@ -50,10 +64,7 @@ def app():
50
  model_path = gr.Dropdown(
51
  label="Model",
52
  choices=[
53
- "gelan-c.pt",
54
- "gelan-e.pt",
55
- "yolov9-c.pt",
56
- "yolov9-e.pt",
57
  ],
58
  value="gelan-e.pt",
59
  )
@@ -95,6 +106,29 @@ def app():
95
  outputs=[output_numpy],
96
  )
97
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98
  gradio_app = gr.Blocks()
99
  with gradio_app:
100
  gr.HTML(
 
4
  from huggingface_hub import hf_hub_download
5
 
6
 
7
+ def attempt_download_from_hub(repo_id, hf_token=None):
8
+ # https://github.com/fcakyon/yolov5-pip/blob/main/yolov5/utils/downloads.py
9
+ from huggingface_hub import hf_hub_download, list_repo_files
10
+ from huggingface_hub.utils._errors import RepositoryNotFoundError
11
+ from huggingface_hub.utils._validators import HFValidationError
12
+ try:
13
+ repo_files = list_repo_files(repo_id=repo_id, repo_type='model', token=hf_token)
14
+ model_file = [f for f in repo_files if f.endswith('.pt')][0]
15
+ file = hf_hub_download(
16
+ repo_id=repo_id,
17
+ filename=model_file,
18
+ repo_type='model',
19
+ token=hf_token,
20
+ )
21
+ return file
22
+ except (RepositoryNotFoundError, HFValidationError):
23
+ return None
24
 
25
 
26
  @spaces.GPU
 
40
  import yolov9
41
 
42
  # Load the model
43
+ model_path = attempt_download_from_hub(model_id)
44
  model = yolov9.load(model_path, device="cuda")
45
 
46
  # Set model parameters
 
64
  model_path = gr.Dropdown(
65
  label="Model",
66
  choices=[
67
+ "kadirnar/yolov9-gelan-c",
 
 
 
68
  ],
69
  value="gelan-e.pt",
70
  )
 
106
  outputs=[output_numpy],
107
  )
108
 
109
+ gr.Examples(
110
+ examples=[
111
+ [
112
+ "data/zidane.jpg",
113
+ "kadirnar/yolov9-gelan-c",
114
+ 640,
115
+ 0.4,
116
+ 0.5,
117
+ ],
118
+ ],
119
+ fn=yolov9_inference,
120
+ inputs=[
121
+ img_path,
122
+ model_path,
123
+ image_size,
124
+ conf_threshold,
125
+ iou_threshold,
126
+ ],
127
+ outputs=[output_numpy],
128
+ cache_examples=True,
129
+ )
130
+
131
+
132
  gradio_app = gr.Blocks()
133
  with gradio_app:
134
  gr.HTML(