Rename app_gradio.py to app.py

app_gradio.py → app.py

@@ -1,24 +1,176 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
+from enum import Enum
+from pathlib import Path
+from typing import List, Union
+
 import gradio as gr
-from PIL import Image
+import numpy as np
+from rapidocr import RapidOCR
+
+
+class InferEngine(Enum):
+    ort = "ONNXRuntime"
+    vino = "OpenVino"
+    paddle = "PaddlePaddle"
+    torch = "PyTorch"
+
+
+def get_ocr_engine(infer_engine: str, lang_det: str, lang_rec: str) -> RapidOCR:
+    if infer_engine == InferEngine.vino.value:
+        return RapidOCR(
+            params={
+                "Global.with_openvino": True,
+                "Global.lang_det": lang_det,
+                "Global.lang_rec": lang_rec,
+            }
+        )
+
+    if infer_engine == InferEngine.paddle.value:
+        return RapidOCR(
+            params={
+                "Global.with_paddle": True,
+                "Global.lang_det": lang_det,
+                "Global.lang_rec": lang_rec,
+            }
+        )
+
+    if infer_engine == InferEngine.torch.value:
+        return RapidOCR(
+            params={
+                "Global.with_torch": True,
+                "Global.lang_det": lang_det,
+                "Global.lang_rec": lang_rec,
+            }
+        )
+
+    return RapidOCR(
+        params={
+            "Global.with_onnx": True,
+            "Global.lang_det": lang_det,
+            "Global.lang_rec": lang_rec,
+        }
+    )
+
+
+def get_ocr_result(
+    img: np.ndarray,
+    text_score,
+    box_thresh,
+    unclip_ratio,
+    lang_det,
+    lang_rec,
+    infer_engine,
+    is_word: str,
+):
+    return_word_box = True if is_word == "Yes" else False
 
-img_path = "images/1.jpg"
+    ocr_engine = get_ocr_engine(infer_engine, lang_det=lang_det, lang_rec=lang_rec)
 
+    ocr_result = ocr_engine(
+        img,
+        text_score=text_score,
+        box_thresh=box_thresh,
+        unclip_ratio=unclip_ratio,
+        return_word_box=return_word_box,
+    )
+    vis_img = ocr_result.vis()
+    if return_word_box:
+        txts, scores, _ = list(zip(*ocr_result.word_results))
+        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))
+    ]
+    return vis_img, ocr_txts, ocr_result.elapse
+
+
+def create_examples() -> List[List[Union[str, float]]]:
+    examples = [
+        [
+            "images/ch_en_num.jpg",
+            0.5,
+            0.5,
+            1.6,
+            "ch_mobile",
+            "ch_mobile",
+            "ONNXRuntime",
+            "No",
+        ],
+        [
+            "images/japan.jpg",
+            0.5,
+            0.5,
+            1.6,
+            "multi_mobile",
+            "japan_mobile",
+            "ONNXRuntime",
+            "No",
+        ],
+        [
+            "images/korean.jpg",
+            0.5,
+            0.5,
+            1.6,
+            "multi_mobile",
+            "korean_mobile",
+            "ONNXRuntime",
+            "No",
+        ],
+        [
+            "images/air_ticket.jpg",
+            0.5,
+            0.5,
+            1.6,
+            "ch_mobile",
+            "ch_mobile",
+            "ONNXRuntime",
+            "No",
+        ],
+        [
+            "images/car_plate.jpeg",
+            0.5,
+            0.5,
+            1.6,
+            "ch_mobile",
+            "ch_mobile",
+            "ONNXRuntime",
+            "No",
+        ],
+        [
+            "images/train_ticket.jpeg",
+            0.5,
+            0.5,
+            1.6,
+            "ch_mobile",
+            "ch_mobile",
+            "ONNXRuntime",
+            "No",
+        ],
+    ]
+    return examples
 
-def test():
-    return Image.open(img_path)
 
+infer_engine_list = [InferEngine[v].value for v in InferEngine.__members__]
 
-example_images = [
-    "images/1.jpg",
-    "images/ch_en_num.jpg",
-    "images/air_ticket.jpg",
-    "images/car_plate.jpeg",
-    "images/train_ticket.jpeg",
-    "images/japan_2.jpg",
-    "images/korean_1.jpg",
+lang_det_list = ["ch_mobile", "ch_server", "en_mobile", "en_server", "multi_mobile"]
+lang_rec_list = [
+    "ch_mobile",
+    "ch_server",
+    "chinese_cht",
+    "en_mobile",
+    "ar_mobile",
+    "cyrillic_mobile",
+    "devanagari_mobile",
+    "japan_mobile",
+    "ka_mobile",
+    "korean_mobile",
+    "latin_mobile",
+    "ta_mobile",
+    "te_mobile",
 ]
 
 custom_css = """
@@ -76,17 +228,62 @@ with gr.Blocks(
             info="控制文本检测框的大小，值越大，检测框整体越大。在出现框截断文字的情况，调大该值。取值范围：[1.5, 2.0]，默认值为1.6",
         )
 
+    with gr.Row():
+        select_infer_engine = gr.Dropdown(
+            choices=infer_engine_list,
+            label="Infer Engine (推理引擎)",
+            value="ONNXRuntime",
+            interactive=True,
+        )
+        lang_det = gr.Dropdown(
+            choices=lang_det_list,
+            label="Det model (文本检测模型)",
+            value=lang_det_list[0],
+            interactive=True,
+        )
+        lang_rec = gr.Dropdown(
+            choices=lang_rec_list,
+            label="Rec model (文本识别模型)",
+            value=lang_rec_list[0],
+            interactive=True,
+        )
+        is_word = gr.Radio(
+            ["Yes", "No"], label="Return word box (返回单字符)", value="No"
+        )
+
     img_input = gr.Image(label="Upload or Select Image", sources="upload")
+
     run_btn = gr.Button("Run")
 
-    run_btn.click(test, inputs=img_input, outputs=gr.Image())
+    img_output = gr.Image(label="Output Image")
+    elapse = gr.Textbox(label="Elapse(s)")
+    ocr_results = gr.Dataframe(
+        label="OCR Txts",
+        headers=["Index", "Txt", "Score"],
+        datatype=["number", "str", "number"],
+        show_copy_button=True,
+    )
+
+    ocr_inputs = [
+        img_input,
+        text_score,
+        box_thresh,
+        unclip_ratio,
+        lang_det,
+        lang_rec,
+        select_infer_engine,
+        is_word,
+    ]
+    run_btn.click(
+        get_ocr_result, inputs=ocr_inputs, outputs=[img_output, ocr_results, elapse]
+    )
 
     examples = gr.Examples(
-        examples=example_images,
-        examples_per_page=len(example_images),
-        inputs=img_input,
-        fn=lambda x: x,  # 简单返回图片路径
-        outputs=img_input,
+        examples=create_examples(),
+        examples_per_page=5,
+        inputs=ocr_inputs,
+        fn=get_ocr_result,
+        outputs=[img_output, ocr_results, elapse],
         cache_examples=False,
     )