Update app.py

app.py

@@ -7,6 +7,7 @@ from typing import List, Union
 
 import gradio as gr
 import numpy as np
+
 from rapidocr import RapidOCR
 
 
@@ -43,13 +44,21 @@ def get_ocr_result(
     lang_rec,
     infer_engine,
     is_word: str,
+    use_module: List[str],
 ):
-    return_word_box = True if is_word == "Yes" else False
+    return_word_box = "Yes" in is_word
+
+    use_det = "use_det" in use_module
+    use_cls = "use_cls" in use_module
+    use_rec = "use_rec" in use_module
 
     ocr_engine = get_ocr_engine(infer_engine, lang_det=lang_det, lang_rec=lang_rec)
 
     ocr_result = ocr_engine(
         img,
+        use_det=use_det,
+        use_cls=use_cls,
+        use_rec=use_rec,
         text_score=text_score,
         box_thresh=box_thresh,
         unclip_ratio=unclip_ratio,
@@ -61,10 +70,19 @@ def get_ocr_result(
         ocr_txts = [[i, txt, score] for i, (txt, score) in enumerate(zip(txts, scores))]
         return vis_img, ocr_txts, ocr_result.elapse
 
-    ocr_txts = [
-        [i, txt, score]
-        for i, (txt, score) in enumerate(zip(ocr_result.txts, ocr_result.scores))
-    ]
+    if use_det and use_rec:
+        ocr_txts = [
+            [i, txt, score]
+            for i, (txt, score) in enumerate(zip(ocr_result.txts, ocr_result.scores))
+        ]
+    elif use_det and not use_cls and not use_rec:
+        ocr_txts = []
+    else not use_det and not use_cls and use_rec:
+        ocr_txts = [
+            [i, txt, score]
+            for i, (txt, score) in enumerate(zip(ocr_result.txts, ocr_result.scores))
+        ]
+
     return vis_img, ocr_txts, ocr_result.elapse
 
 
@@ -128,7 +146,7 @@ with gr.Blocks(
     gr.HTML(
         """
         <h1 style='text-align: center;font-size:40px'>Rapid⚡OCR</h1>
-        
+
         <div style="display: flex; justify-content: center; gap: 10px;">
             <a href=""><img src="https://img.shields.io/badge/Python->=3.6-aff.svg"></a>
             <a href="https://rapidai.github.io/RapidOCRDocs"><img src="https://img.shields.io/badge/Docs-link-aff.svg"></a>
@@ -166,6 +184,12 @@ with gr.Blocks(
         )
 
     with gr.Row():
+        use_module = gr.CheckboxGroup(
+            ["use_det", "use_cls", "use_rec"],
+            label="Use module (使用哪些模块)",
+            value=["use_det", "use_cls", "use_rec"],
+            interactive=True,
+        )
         select_infer_engine = gr.Dropdown(
             choices=infer_engine_list,
             label="Infer Engine (推理引擎)",
@@ -210,6 +234,7 @@ with gr.Blocks(
         lang_rec,
         select_infer_engine,
         is_word,
+        use_module,
     ]
     run_btn.click(
         get_ocr_result, inputs=ocr_inputs, outputs=[img_output, ocr_results, elapse]