dev(narugo): add more usable model

app.py

@@ -124,11 +124,12 @@ def load_classes() -> List[str]:
         return json.load(f)
 
 
-def get_tags_from_image(pic: Image.Image, threshold: float = 0.7, size: int = 512, keep_ratio: bool = False):
+def get_tags_from_image(pic: Image.Image, threshold: float = 0.7, size: int = 512, keep_ratio: bool = False,
+                        model_name='TResnet-D-FLq_ema_6-10000.onnx'):
     real_input = image_to_tensor(pic, size, keep_ratio)
     real_input = real_input.reshape(1, *real_input.shape)
 
-    model = _open_onnx_model(get_onnx_model_file(), get_onnx_provider('cpu'))
+    model = _open_onnx_model(get_onnx_model_file(model_name), get_onnx_provider('cpu'))
     native_output, = model.run(['output'], {'input': real_input})
 
     output = (1 / (1 + np.exp(-native_output))).reshape(-1)
@@ -140,10 +141,10 @@ def get_tags_from_image(pic: Image.Image, threshold: float = 0.7, size: int = 51
 RE_SPECIAL = re.compile(r'([\\()])')
 
 
-def image_to_mldanbooru_tags(pic: Image.Image, threshold: float, size: int, keep_ratio: bool,
+def image_to_mldanbooru_tags(pic: Image.Image, threshold: float, size: int, keep_ratio: bool, model: str,
                              use_spaces: bool, use_escape: bool, include_ranks: bool, score_descend: bool) \
         -> Tuple[str, Mapping[str, float]]:
-    filtered_tags = get_tags_from_image(pic, threshold, size, keep_ratio)
+    filtered_tags = get_tags_from_image(pic, threshold, size, keep_ratio, model)
 
     text_items = []
     tags_pairs = filtered_tags.items()
@@ -163,6 +164,12 @@ def image_to_mldanbooru_tags(pic: Image.Image, threshold: float, size: int, keep
     return output_text, filtered_tags
 
 
+MODELS = [
+    'TResnet-D-FLq_ema_6-10000.onnx',
+    'TResnet-D-FLq_ema_4-10000.onnx',
+    'TResnet-D-FLq_ema_2-40000.onnx',
+]
+
 if __name__ == '__main__':
     with gr.Blocks() as demo:
         with gr.Row():
@@ -172,6 +179,8 @@ if __name__ == '__main__':
                     gr_threshold = gr.Slider(0.0, 1.0, 0.7, label='Tagging Confidence Threshold')
                     gr_image_size = gr.Slider(128, 960, 512, step=32, label='Image for Recognition')
                     gr_keep_ratio = gr.Checkbox(value=False, label='Keep the Ratio')
+                with gr.Row():
+                    gr_model = gr.Dropdown(MODELS, value=MODELS[0], label='Model')
                 with gr.Row():
                     gr_space = gr.Checkbox(value=False, label='Use Space Instead Of _')
                     gr_escape = gr.Checkbox(value=True, label='Use Text Escape')
@@ -190,7 +199,8 @@ if __name__ == '__main__':
         gr_btn_submit.click(
             image_to_mldanbooru_tags,
             inputs=[
-                gr_input_image, gr_threshold, gr_image_size, gr_keep_ratio,
+                gr_input_image, gr_threshold, gr_image_size,
+                gr_keep_ratio, gr_model,
                 gr_space, gr_escape, gr_confidence, gr_order
             ],
             outputs=[gr_output_text, gr_tags],