Update app.py

app.py

@@ -1,9 +1,13 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
+from pathlib import Path
 from typing import List, Union
 
 import gradio as gr
+from omegaconf import OmegaConf
+
+import rapidocr
 from rapidocr import (
     EngineType,
     LangCls,
@@ -191,6 +195,66 @@ def create_examples() -> List[List[Union[str, float]]]:
     return examples
 
 
+def export_yaml(
+    img_input,
+    text_score,
+    box_thresh,
+    unclip_ratio,
+    det_engine,
+    lang_det,
+    det_model_type,
+    det_ocr_version,
+    cls_engine,
+    lang_cls,
+    cls_model_type,
+    cls_ocr_version,
+    rec_engine,
+    lang_rec,
+    rec_model_type,
+    rec_ocr_version,
+    is_word,
+    use_module,
+):
+    default_yaml_path = Path(rapidocr.__file__).parent / "config.yaml"
+    cfg = OmegaConf.load(default_yaml_path)
+
+    params = {
+        "Global": {
+            "use_det": "use_det" in use_module,
+            "use_cls": "use_cls" in use_module,
+            "use_rec": "use_rec" in use_module,
+            "return_word_box": "Yes" in is_word,
+            "text_score": text_score,
+            "box_thresh": box_thresh,
+        },
+        "Det": {
+            "engine_type": det_engine,
+            "lang_type": lang_det,
+            "model_type": det_model_type,
+            "ocr_version": det_ocr_version,
+            "box_thresh": box_thresh,
+            "unclip_ratio": unclip_ratio,
+        },
+        "Cls": {
+            "engine_type": cls_engine,
+            "lang_type": lang_cls,
+            "model_type": cls_model_type,
+            "ocr_version": cls_ocr_version,
+        },
+        "Rec": {
+            "engine_type": rec_engine,
+            "lang_type": lang_rec,
+            "model_type": rec_model_type,
+            "ocr_version": rec_ocr_version,
+        },
+    }
+    cfg = OmegaConf.merge(cfg, params)
+
+    save_path = Path(__file__).resolve().parent / "config.yaml"
+    OmegaConf.save(cfg, save_path)
+    return save_path
+
+
 custom_css = """
     body {font-family: body {font-family: 'Helvetica Neue', Helvetica;}
     .gr-button {background-color: #4CAF50; color: white; border: none; padding: 10px 20px; border-radius: 5px;}
@@ -222,7 +286,7 @@ with gr.Blocks(
     )
     img_input = gr.Image(label="Upload or Select Image", sources="upload")
 
-    with gr.Accordion("Advanced Setting", open=False):
+    with gr.Accordion("Parameter Setting", open=False):
         with gr.Row():
             text_score = gr.Slider(
                 label="text_score",
@@ -345,9 +409,12 @@ with gr.Blocks(
                 ["Yes", "No"], label="Return word box (返回单字符)", value="No"
             )
 
-    with gr.Row():
         btn_export_cfg = gr.Button("Export Config YAML")
-        run_btn = gr.Button("Run")
+        download_btn_hidden = gr.DownloadButton(
+            visible=False, elem_id="download_btn_hidden"
+        )
+
+    run_btn = gr.Button("Run")
 
     img_output = gr.Image(label="Output Image")
     elapse = gr.Textbox(label="Elapse(s)")
@@ -382,6 +449,15 @@ with gr.Blocks(
         get_ocr_result, inputs=ocr_inputs, outputs=[img_output, ocr_results, elapse]
     )
 
+    btn_export_cfg.click(
+        fn=export_yaml, inputs=ocr_inputs, outputs=[download_btn_hidden]
+    ).then(
+        fn=None,
+        inputs=None,
+        outputs=None,
+        js="() => document.querySelector('#download_btn_hidden').click()",
+    )
+
     examples = gr.Examples(
         examples=create_examples(),
         examples_per_page=5,