# -*- 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( """

Rapid⚡OCRv3

PyPI
""" ) 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)