narugo1992 commited on
Commit
726a752
1 Parent(s): 94b61d2

dev(narugo): add new model

Browse files
Files changed (1) hide show
  1. app.py +12 -9
app.py CHANGED
@@ -97,7 +97,16 @@ def image_to_tensor(pic: Image.Image, size: int = 512, keep_ratio: float = True,
97
  return to_tensor(resize(fill_background(pic, background), size, keep_ratio))
98
 
99
 
100
- def get_onnx_model_file(name='TResnet-D-FLq_ema_6-10000.onnx'):
 
 
 
 
 
 
 
 
 
101
  return hf_hub_download(
102
  repo_id='deepghs/ml-danbooru-onnx',
103
  filename=name,
@@ -125,7 +134,7 @@ def load_classes() -> List[str]:
125
 
126
 
127
  def get_tags_from_image(pic: Image.Image, threshold: float = 0.7, size: int = 512, keep_ratio: bool = False,
128
- model_name='TResnet-D-FLq_ema_6-10000.onnx'):
129
  real_input = image_to_tensor(pic, size, keep_ratio)
130
  real_input = real_input.reshape(1, *real_input.shape)
131
 
@@ -164,12 +173,6 @@ def image_to_mldanbooru_tags(pic: Image.Image, threshold: float, size: int, keep
164
  return output_text, filtered_tags
165
 
166
 
167
- MODELS = [
168
- 'TResnet-D-FLq_ema_6-10000.onnx',
169
- 'TResnet-D-FLq_ema_4-10000.onnx',
170
- 'TResnet-D-FLq_ema_2-40000.onnx',
171
- ]
172
-
173
  if __name__ == '__main__':
174
  with gr.Blocks() as demo:
175
  with gr.Row():
@@ -180,7 +183,7 @@ if __name__ == '__main__':
180
  gr_image_size = gr.Slider(128, 960, 640, step=32, label='Image for Recognition')
181
  gr_keep_ratio = gr.Checkbox(value=False, label='Keep the Ratio')
182
  with gr.Row():
183
- gr_model = gr.Dropdown(MODELS, value=MODELS[0], label='Model')
184
  with gr.Row():
185
  gr_space = gr.Checkbox(value=False, label='Use Space Instead Of _')
186
  gr_escape = gr.Checkbox(value=True, label='Use Text Escape')
 
97
  return to_tensor(resize(fill_background(pic, background), size, keep_ratio))
98
 
99
 
100
+ MODELS = [
101
+ 'TResnet-D-FLq_ema_6-30000.onnx',
102
+ 'TResnet-D-FLq_ema_6-10000.onnx',
103
+ 'TResnet-D-FLq_ema_4-10000.onnx',
104
+ 'TResnet-D-FLq_ema_2-40000.onnx',
105
+ ]
106
+ DEFAULT_MODEL = MODELS[0]
107
+
108
+
109
+ def get_onnx_model_file(name=DEFAULT_MODEL):
110
  return hf_hub_download(
111
  repo_id='deepghs/ml-danbooru-onnx',
112
  filename=name,
 
134
 
135
 
136
  def get_tags_from_image(pic: Image.Image, threshold: float = 0.7, size: int = 512, keep_ratio: bool = False,
137
+ model_name=DEFAULT_MODEL):
138
  real_input = image_to_tensor(pic, size, keep_ratio)
139
  real_input = real_input.reshape(1, *real_input.shape)
140
 
 
173
  return output_text, filtered_tags
174
 
175
 
 
 
 
 
 
 
176
  if __name__ == '__main__':
177
  with gr.Blocks() as demo:
178
  with gr.Row():
 
183
  gr_image_size = gr.Slider(128, 960, 640, step=32, label='Image for Recognition')
184
  gr_keep_ratio = gr.Checkbox(value=False, label='Keep the Ratio')
185
  with gr.Row():
186
+ gr_model = gr.Dropdown(MODELS, value=DEFAULT_MODEL, label='Model')
187
  with gr.Row():
188
  gr_space = gr.Checkbox(value=False, label='Use Space Instead Of _')
189
  gr_escape = gr.Checkbox(value=True, label='Use Text Escape')