app.py

# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Sequence, Tuple

import gradio as gr
import rapidocr
from omegaconf import OmegaConf
from rapidocr import (
    EngineType,
    LangCls,
    LangDet,
    LangRec,
    ModelType,
    OCRVersion,
    RapidOCR,
)

OCR_INPUT_FIELDS: Tuple[str, ...] = (
    "img_input",
    "text_score",
    "box_thresh",
    "unclip_ratio",
    "max_side_len",
    "limit_side_len",
    "limit_type",
    "use_dilation",
    "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_INPUT_VALUES: Dict[str, Any] = {
    "img_input": None,
    "text_score": 0.5,
    "box_thresh": 0.5,
    "unclip_ratio": 1.6,
    "max_side_len": 2000,
    "limit_side_len": 736,
    "limit_type": "min",
    "use_dilation": "True",
    "det_engine": EngineType.ONNXRUNTIME.value,
    "lang_det": LangDet.CH.value,
    "det_model_type": ModelType.SMALL.value,
    "det_ocr_version": OCRVersion.PPOCRV6.value,
    "cls_engine": EngineType.ONNXRUNTIME.value,
    "lang_cls": LangCls.CH.value,
    "cls_model_type": ModelType.MOBILE.value,
    "cls_ocr_version": OCRVersion.PPOCRV4.value,
    "rec_engine": EngineType.ONNXRUNTIME.value,
    "lang_rec": LangRec.CH.value,
    "rec_model_type": ModelType.SMALL.value,
    "rec_ocr_version": OCRVersion.PPOCRV6.value,
    "is_word": "No",
    "use_module": ["use_det", "use_cls", "use_rec"],
}


@dataclass(frozen=True)
class OCRAppConfig:
    img_input: Any
    text_score: float
    box_thresh: float
    unclip_ratio: float
    max_side_len: int
    limit_side_len: int
    limit_type: str
    use_dilation: str
    det_engine: str
    lang_det: str
    det_model_type: str
    det_ocr_version: str
    cls_engine: str
    lang_cls: str
    cls_model_type: str
    cls_ocr_version: str
    rec_engine: str
    lang_rec: str
    rec_model_type: str
    rec_ocr_version: str
    is_word: str
    use_module: Sequence[str]

    @classmethod
    def from_values(cls, values: Sequence[Any]) -> "OCRAppConfig":
        if len(values) != len(OCR_INPUT_FIELDS):
            raise ValueError(
                f"参数数量不匹配：需要 {len(OCR_INPUT_FIELDS)} 个，实际收到 {len(values)} 个"
            )

        raw_config = dict(zip(OCR_INPUT_FIELDS, values))
        for key in ("text_score", "box_thresh", "unclip_ratio"):
            raw_config[key] = float(raw_config[key])
        for key in ("max_side_len", "limit_side_len"):
            raw_config[key] = int(raw_config[key])
        return cls(**raw_config)

    @property
    def selected_modules(self) -> Sequence[str]:
        return self.use_module or []

    @property
    def return_word_box(self) -> bool:
        return self.is_word == "Yes"

    @property
    def use_det(self) -> bool:
        return "use_det" in self.selected_modules

    @property
    def use_cls(self) -> bool:
        return "use_cls" in self.selected_modules

    @property
    def use_rec(self) -> bool:
        return "use_rec" in self.selected_modules

    @property
    def use_dilation_bool(self) -> bool:
        return self.use_dilation == "True"

    def to_rapidocr_params(self) -> Dict[str, Any]:
        return {
            "Global.max_side_len": self.max_side_len,
            "Det.engine_type": EngineType(self.det_engine),
            "Det.lang_type": LangDet(self.lang_det),
            "Det.model_type": ModelType(self.det_model_type),
            "Det.ocr_version": OCRVersion(self.det_ocr_version),
            "Det.use_dilation": self.use_dilation_bool,
            "Det.limit_side_len": self.limit_side_len,
            "Det.limit_type": self.limit_type,
            "Cls.engine_type": EngineType(self.cls_engine),
            "Cls.lang_type": LangCls(self.lang_cls),
            "Cls.model_type": ModelType(self.cls_model_type),
            "Cls.ocr_version": OCRVersion(self.cls_ocr_version),
            "Rec.engine_type": EngineType(self.rec_engine),
            "Rec.lang_type": LangRec(self.lang_rec),
            "Rec.model_type": ModelType(self.rec_model_type),
            "Rec.ocr_version": OCRVersion(self.rec_ocr_version),
        }

    def to_yaml_params(self) -> Dict[str, Any]:
        return {
            "Global": {
                "max_side_len": self.max_side_len,
                "use_det": self.use_det,
                "use_cls": self.use_cls,
                "use_rec": self.use_rec,
                "return_word_box": self.return_word_box,
                "text_score": self.text_score,
                "box_thresh": self.box_thresh,
            },
            "Det": {
                "engine_type": self.det_engine,
                "lang_type": self.lang_det,
                "model_type": self.det_model_type,
                "ocr_version": self.det_ocr_version,
                "box_thresh": self.box_thresh,
                "unclip_ratio": self.unclip_ratio,
                "use_dilation": self.use_dilation_bool,
                "limit_side_len": self.limit_side_len,
                "limit_type": self.limit_type,
            },
            "Cls": {
                "engine_type": self.cls_engine,
                "lang_type": self.lang_cls,
                "model_type": self.cls_model_type,
                "ocr_version": self.cls_ocr_version,
            },
            "Rec": {
                "engine_type": self.rec_engine,
                "lang_type": self.lang_rec,
                "model_type": self.rec_model_type,
                "ocr_version": self.rec_ocr_version,
            },
        }


def _build_config(values: Sequence[Any]) -> OCRAppConfig:
    return OCRAppConfig.from_values(values)


def _build_example(**overrides: Any) -> List[Any]:
    example = {**DEFAULT_INPUT_VALUES, **overrides}
    return [example[field] for field in OCR_INPUT_FIELDS]


def _format_ocr_result(ocr_result, config: OCRAppConfig):
    vis_img = ocr_result.vis()
    if config.return_word_box:
        full_word_results = [
            word_result
            for line_word_results in (ocr_result.word_results or [])
            for word_result in line_word_results
        ]
        ocr_txts = [
            [idx, txt, score] for idx, (txt, score, _) in enumerate(full_word_results)
        ]
        return vis_img, ocr_txts, ocr_result.elapse

    if not config.use_rec:
        return vis_img, [], ocr_result.elapse

    txts = ocr_result.txts or []
    scores = ocr_result.scores or []
    ocr_txts = [[idx, txt, score] for idx, (txt, score) in enumerate(zip(txts, scores))]
    return vis_img, ocr_txts, ocr_result.elapse


def get_ocr_result(*values):
    try:
        config = _build_config(values)
        ocr_engine = RapidOCR(params=config.to_rapidocr_params())

        ocr_result = ocr_engine(
            config.img_input,
            use_det=config.use_det,
            use_cls=config.use_cls,
            use_rec=config.use_rec,
            text_score=config.text_score,
            box_thresh=config.box_thresh,
            unclip_ratio=config.unclip_ratio,
            return_word_box=config.return_word_box,
        )
    except Exception as e:
        err_msg = f"模型加载/识别失败：{str(e)}，详细参见 Logs"
        gr.Warning(err_msg)
        print(err_msg)
        return None, [], 0.0

    return _format_ocr_result(ocr_result, config)


def create_examples() -> List[List[Any]]:
    return [
        _build_example(img_input="images/multi.jpg"),
        _build_example(img_input="images/ch_en_num.jpg"),
        _build_example(img_input="images/hand_writen.jpeg"),
        _build_example(img_input="images/japan.jpg", lang_rec=LangRec.JAPAN.value),
        _build_example(
            img_input="images/korean.jpg",
            det_model_type=ModelType.MOBILE.value,
            det_ocr_version=OCRVersion.PPOCRV5.value,
            lang_rec=LangRec.KOREAN.value,
            rec_model_type=ModelType.MOBILE.value,
            rec_ocr_version=OCRVersion.PPOCRV5.value,
        ),
    ]


def export_yaml(*values):
    config = _build_config(values)
    default_yaml_path = Path(rapidocr.__file__).parent / "config.yaml"
    cfg = OmegaConf.load(default_yaml_path)
    cfg = OmegaConf.merge(cfg, config.to_yaml_params())

    save_path = Path(__file__).resolve().parent / "config.yaml"
    OmegaConf.save(cfg, save_path)
    return save_path


custom_css = """
    body {font-family: 'Helvetica Neue', Helvetica;}
    .gr-button {background-color: #4CAF50; color: white; border: none; padding: 10px 20px; border-radius: 5px;}
    .gr-button:hover {background-color: #45a049;}
    .gr-textbox {margin-bottom: 15px;}
    .example-button {background-color: #1E90FF; color: white; border: none; padding: 8px 15px; border-radius: 5px; margin: 5px;}
    .example-button:hover {background-color: #FF4500;}
    .tall-radio .gr-radio-item {padding: 15px 0; min-height: 50px; display: flex; align-items: center;}
    .tall-radio label {font-size: 16px;}
    .output-image, .input-image, .image-preview {height: 300px !important}
"""

with gr.Blocks(title="Rapid⚡OCR Demo", css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <h1 style='text-align: center;font-size:40px'>Rapid⚡OCRv3</h1>

        <div style="display: flex; justify-content: center; gap: 10px;">
            <a href=""><img src="https://img.shields.io/badge/Python->=3.8-aff.svg"></a>
            <a href="https://rapidai.github.io/RapidOCRDocs"><img src="https://img.shields.io/badge/Docs-link-aff.svg"></a>
            <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
            <a href="https://pepy.tech/project/rapidocr"><img src="https://static.pepy.tech/personalized-badge/rapidocr?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads%20rapidocr"></a>
            <a href="https://pypi.org/project/rapidocr/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapidocr"></a>
            <a href="https://github.com/RapidAI/RapidOCR"><img src="https://img.shields.io/github/stars/RapidAI/RapidOCR?color=ccf"></a>
        </div>
    """
    )
    img_input = gr.Image(label="Upload or Select Image", sources="upload")

    with gr.Accordion("Parameter Setting", open=False):
        with gr.Row():
            text_score = gr.Slider(
                label="text_score",
                minimum=0,
                maximum=1.0,
                value=0.5,
                step=0.1,
                info="文本识别结果是正确的置信度，值越大，显示出的识别结果更准确。存在漏检时，调低该值。取值范围：[0, 1.0]，默认值为0.5",
            )
            box_thresh = gr.Slider(
                label="box_thresh",
                minimum=0,
                maximum=1.0,
                value=0.5,
                step=0.1,
                info="检测到的框是文本的概率，值越大，框中是文本的概率就越大。存在漏检时，调低该值。取值范围：[0, 1.0]，默认值为0.5",
            )
            unclip_ratio = gr.Slider(
                label="unclip_ratio",
                minimum=1.5,
                maximum=2.0,
                value=1.6,
                step=0.1,
                info="控制文本检测框的大小，值越大，检测框整体越大。在出现框截断文字的情况，调大该值。取值范围：[1.5, 2.0]，默认值为1.6",
            )
            max_side_len = gr.Number(
                value=2000,
                label="max_side_len",
                info="如果输入图像的最大边大于`max_side_len`，则会按宽高比，将最大边缩放到`max_side_len`。默认为2000px",
                interactive=True,
                minimum=20,
            )
            limit_side_len = gr.Number(
                value=736,
                label="limit_side_len",
                info="如果输入图像的最大边大于`limit_side_len`，则会按宽高比，将最大边缩放到`limit_side_len`。默认为736px",
                interactive=True,
                minimum=20,
            )
            limit_type = gr.Radio(
                ["min", "max"],
                label="limit_type",
                value="min",
                info="缩放图像时，按最小边还是最大边进行缩放。默认为min",
                interactive=True,
            )
            use_dilation = gr.Radio(
                ["True", "False"],
                label="use_dilation",
                value="True",
                info="是否使用膨胀操作，膨胀操作可以让检测框更大，避免出现框截断文字的情况。默认为True",
                interactive=True,
            )
        with gr.Row():
            with gr.Row():
                gr.Markdown("Det")
                det_engine = gr.Dropdown(
                    choices=[v.value for v in EngineType],
                    label="EngineType",
                    value=EngineType.ONNXRUNTIME.value,
                    interactive=True,
                    scale=0,
                    allow_custom_value=True,
                )
                lang_det = gr.Dropdown(
                    choices=[v.value for v in LangDet],
                    label="LangDet",
                    value=LangDet.CH.value,
                    interactive=True,
                    scale=1,
                    allow_custom_value=True,
                )
                det_model_type = gr.Dropdown(
                    choices=[v.value for v in ModelType],
                    label="ModelType",
                    value=ModelType.SMALL.value,
                    interactive=True,
                    scale=1,
                    allow_custom_value=True,
                )
                det_ocr_version = gr.Dropdown(
                    choices=[v.value for v in OCRVersion],
                    label="OCR Version",
                    value=OCRVersion.PPOCRV6.value,
                    interactive=True,
                    scale=1,
                    allow_custom_value=True,
                )
            with gr.Row():
                gr.Markdown("Cls")
                cls_engine = gr.Dropdown(
                    choices=[v.value for v in EngineType],
                    label="EngineType",
                    value=EngineType.ONNXRUNTIME.value,
                    interactive=True,
                    allow_custom_value=True,
                )
                lang_cls = gr.Dropdown(
                    choices=[v.value for v in LangCls],
                    label="LangCls",
                    value=LangCls.CH.value,
                    interactive=True,
                    allow_custom_value=True,
                )
                cls_model_type = gr.Dropdown(
                    choices=[v.value for v in ModelType],
                    label="ModelType",
                    value=ModelType.MOBILE.value,
                    interactive=True,
                    allow_custom_value=True,
                )
                cls_ocr_version = gr.Dropdown(
                    choices=[v.value for v in OCRVersion],
                    label="OCR Version",
                    value=OCRVersion.PPOCRV4.value,
                    interactive=True,
                    allow_custom_value=True,
                )
            with gr.Row():
                gr.Markdown("Rec")
                rec_engine = gr.Dropdown(
                    choices=[v.value for v in EngineType],
                    label="EngineType",
                    value=EngineType.ONNXRUNTIME.value,
                    interactive=True,
                    allow_custom_value=True,
                )
                lang_rec = gr.Dropdown(
                    choices=[v.value for v in LangRec],
                    label="LangRec",
                    value=LangRec.CH.value,
                    interactive=True,
                    allow_custom_value=True,
                )
                rec_model_type = gr.Dropdown(
                    choices=[v.value for v in ModelType],
                    label="ModelType",
                    value=ModelType.SMALL.value,
                    interactive=True,
                    allow_custom_value=True,
                )
                rec_ocr_version = gr.Dropdown(
                    choices=[v.value for v in OCRVersion],
                    label="OCR Version",
                    value=OCRVersion.PPOCRV6.value,
                    interactive=True,
                    allow_custom_value=True,
                )

        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,
            )

            is_word = gr.Radio(
                ["Yes", "No"], label="Return word box (返回单字符)", value="No"
            )

    with gr.Row():
        run_btn = gr.Button("Run")
        btn_export_cfg = gr.Button("Export Config YAML")
        download_btn_hidden = gr.DownloadButton(
            visible=False, elem_id="download_btn_hidden"
        )

    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,
    )

    input_components = locals()
    # 组件变量名与 OCR_INPUT_FIELDS 保持一致，避免多处维护参数顺序。
    ocr_inputs = [input_components[field] for field in OCR_INPUT_FIELDS]
    run_btn.click(
        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,
        inputs=ocr_inputs,
        fn=get_ocr_result,
        outputs=[img_output, ocr_results, elapse],
        cache_examples=False,
    )


if __name__ == "__main__":
    demo.launch(debug=True)