Spaces:
Sleeping
Sleeping
narugo1992
commited on
Commit
•
726a752
1
Parent(s):
94b61d2
dev(narugo): add new model
Browse files
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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=
|
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=
|
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')
|