diff --git a/python/demo.py b/python/demo.py index b7ff421..aa92dd8 100644 --- a/python/demo.py +++ b/python/demo.py @@ -18,7 +18,7 @@ from rapidocr import RapidOCR, VisRes # params={"Global.with_onnx": True, "EngineConfig.onnxruntime.use_cuda": True} # ) # engine = RapidOCR(params={"Global.with_torch": True, "Global.lang": "ch"}) -engine = RapidOCR(params={"Global.with_openvino": True, "Global.lang": "ch"}) +engine = RapidOCR(params={"Global.with_openvino": True, "Global.lang_rec": "ch_server"}) vis = VisRes() image_path = "tests/test_files/ch_en_num.jpg" diff --git a/python/rapidocr/config.yaml b/python/rapidocr/config.yaml index 448d9b6..7ef08b1 100644 --- a/python/rapidocr/config.yaml +++ b/python/rapidocr/config.yaml @@ -1,14 +1,12 @@ Global: - lang: "ch" + lang_det: "ch_mobile" # ch_server + lang_rec: "ch_mobile" text_score: 0.5 use_det: true use_cls: true use_rec: true - use_default_mobile: false - use_default_server: false - min_height: 30 width_height_ratio: 8 max_side_len: 2000 diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index be84bb1..c524820 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/python/rapidocr/inference_engine/base.py @@ -103,9 +103,19 @@ class InferSession(abc.ABC): @classmethod def get_model_url(cls, engine_name: str, task_type: str, lang: str) -> str: + lang, model_type = lang.split("_") model_dict = cls.model_info[engine_name]["PP-OCRv4"][task_type] + + # 优先查找 server 模型 + if model_type == "server": + for k in model_dict: + prefix = k.split("_")[0] + if lang == prefix and "server" in k: + return model_dict[k] + for k in model_dict: prefix = k.split("_")[0] if lang == prefix: return model_dict[k] + raise KeyError("Model not found") diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index 46068ba..acc0157 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -16,11 +16,11 @@ from .ch_ppocr_rec import TextRecInput, TextRecognizer, TextRecOutput from .inference_engine.base import get_engine_name from .utils import ( LoadImage, + ParseLang, RapidOCROutput, VisRes, add_round_letterbox, increase_min_side, - parse_lang, reduce_max_side, ) from .utils.parse_parameters import ParseParams @@ -43,28 +43,28 @@ class RapidOCR: engine_name = get_engine_name(config) - det_lang, rec_lang = parse_lang(config.Global.lang) + lang_det, lang_rec = ParseLang()(config.Global.lang_det, config.Global.lang_rec) self.text_score = config.Global.text_score self.min_height = config.Global.min_height self.width_height_ratio = config.Global.width_height_ratio self.use_det = config.Global.use_det - config.Det.lang = det_lang + config.Det.lang = lang_det config.Det.engine_name = engine_name config.Det.engine_cfg = config.EngineConfig[engine_name] config.Det.task_type = "det" self.text_det = TextDetector(config.Det) self.use_cls = config.Global.use_cls - config.Cls.lang = "ch" + config.Cls.lang = "ch_mobile" config.Cls.engine_name = engine_name config.Cls.engine_cfg = config.EngineConfig[engine_name] config.Cls.task_type = "cls" self.text_cls = TextClassifier(config.Cls) self.use_rec = config.Global.use_rec - config.Rec.lang = rec_lang + config.Rec.lang = lang_rec config.Rec.engine_name = engine_name config.Rec.engine_cfg = config.EngineConfig[engine_name] config.Rec.task_type = "rec" diff --git a/python/rapidocr/utils/__init__.py b/python/rapidocr/utils/__init__.py index 3282700..dcebccf 100644 --- a/python/rapidocr/utils/__init__.py +++ b/python/rapidocr/utils/__init__.py @@ -3,7 +3,7 @@ # @Contact: liekkaskono@163.com from .load_image import LoadImage, LoadImageError from .logger import Logger -from .parse_parameters import ParseParams, parse_lang +from .parse_parameters import ParseLang, ParseParams from .process_img import add_round_letterbox, increase_min_side, reduce_max_side from .typings import RapidOCROutput from .utils import download_file diff --git a/python/rapidocr/utils/parse_parameters.py b/python/rapidocr/utils/parse_parameters.py index 8bece79..ba9cdb3 100644 --- a/python/rapidocr/utils/parse_parameters.py +++ b/python/rapidocr/utils/parse_parameters.py @@ -2,7 +2,7 @@ # @Author: SWHL # @Contact: liekkaskono@163.com from pathlib import Path -from typing import Dict, Union +from typing import Dict, Tuple, Union import numpy as np from omegaconf import DictConfig, OmegaConf @@ -12,6 +12,121 @@ root_dir = Path(__file__).resolve().parent.parent InputType = Union[str, np.ndarray, bytes, Path, Image.Image] +class ParseLang: + def __init__(self): + self.latin_lang = [ + "af", + "az", + "bs", + "cs", + "cy", + "da", + "de", + "es", + "et", + "fr", + "ga", + "hr", + "hu", + "id", + "is", + "it", + "ku", + "la", + "lt", + "lv", + "mi", + "ms", + "mt", + "nl", + "no", + "oc", + "pi", + "pl", + "pt", + "ro", + "rs_latin", + "sk", + "sl", + "sq", + "sv", + "sw", + "tl", + "tr", + "uz", + "vi", + "french", + "german", + ] + + self.arabic_lang = ["ar", "fa", "ug", "ur"] + self.cyrillic_lang = [ + "ru", + "rs_cyrillic", + "be", + "bg", + "uk", + "mn", + "abq", + "ady", + "kbd", + "ava", + "dar", + "inh", + "che", + "lbe", + "lez", + "tab", + ] + self.devanagari_lang = [ + "hi", + "mr", + "ne", + "bh", + "mai", + "ang", + "bho", + "mah", + "sck", + "new", + "gom", + "sa", + "bgc", + ] + + def __call__(self, lang_det: str, lang_rec: str) -> Tuple[str, str]: + lang_det, det_model_type = lang_det.split("_") + lang_rec, rec_model_type = lang_rec.split("_") + return f"{lang_det}_{det_model_type}", f"{lang_rec}_{rec_model_type}" + + def parse_det_lang(self, lang: str) -> str: + if lang == "ch": + return "ch" + + if lang in ["en", "latin"]: + return "en" + + return "Multilingual" + + def parse_rec_lang(self, lang: str): + if lang in self.latin_lang: + return "latin" + + if lang in self.arabic_lang: + return "arabic" + + if lang in self.cyrillic_lang: + return "cyrillic" + + if lang in self.devanagari_lang: + return "devanagari" + + if lang in ["ch", "en"]: + return + + raise ValueError(f"lang: {lang} is not in the supported list.") + + def parse_lang(lang): latin_lang = [ "af", diff --git a/python/tests/test_main.py b/python/tests/test_main.py index 8029306..6565c77 100644 --- a/python/tests/test_main.py +++ b/python/tests/test_main.py @@ -38,6 +38,13 @@ def get_engine(params: Optional[Dict[str, Any]] = None): return engine +def test_server_rec(): + engine = RapidOCR(params={"Global.lang_rec": "ch_server"}) + result = engine(img_path) + assert result.txts is not None + assert result.txts[0] == "正品促销" + + @pytest.mark.parametrize("cmd,gt", [(f"--img_path {img_path}", "正品促销")]) def test_cli(capsys, cmd, gt): main(shlex.split(cmd))