diff --git a/python/rapidocr/ch_ppocr_rec/main.py b/python/rapidocr/ch_ppocr_rec/main.py index d49ad4d..e902717 100644 --- a/python/rapidocr/ch_ppocr_rec/main.py +++ b/python/rapidocr/ch_ppocr_rec/main.py @@ -23,7 +23,7 @@ from rapidocr.inference_engine.base import FileInfo, get_engine from ..utils.download_file import DownloadFile, DownloadFileInput from ..utils.log import logger -from ..utils.typings import LangRec +from ..utils.model_resolver import normalize_lang from ..utils.utils import reorder_bidi_for_display from ..utils.vis_res import VisRes from .typings import TextRecInput, TextRecOutput @@ -50,6 +50,8 @@ class TextRecognizer: self.cfg = cfg + self.RTL_LANGS = {"arabic", "fa", "ur"} + def get_character_dict(self, cfg): character = None dict_path = cfg.get("rec_keys_path", None) @@ -137,7 +139,7 @@ class TextRecognizer: all_line_results, all_word_results = list(zip(*rec_res)) txts, scores = list(zip(*all_line_results)) - if self.cfg.lang_type == LangRec.ARABIC: + if normalize_lang(self.cfg.lang_type) in self.RTL_LANGS: txts = reorder_bidi_for_display(txts) elapse = time.perf_counter() - start_time diff --git a/python/rapidocr/default_models.yaml b/python/rapidocr/default_models.yaml index a395dda..996de5e 100644 --- a/python/rapidocr/default_models.yaml +++ b/python/rapidocr/default_models.yaml @@ -117,23 +117,23 @@ onnxruntime: SHA256: 7d3c02ef6c7da8ae08b4347cc7695b2081aae68c325d64375724ecf39c99e743 PP-OCRv6: det: - ch_PP-OCRv6_det_tiny: + multi_PP-OCRv6_det_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/det/PP-OCRv6_det_tiny.onnx SHA256: f42c0fbd294d95eac1a550e131b277dac97462c8025fa4b6c3cec1b7894bd3d5 - ch_PP-OCRv6_det_small: + multi_PP-OCRv6_det_small: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/det/PP-OCRv6_det_small.onnx SHA256: 090f04abcd9d9a7498bc4ebf677e4cb9bdce1fe4197ddb7e529f1ef44e1ff94f - ch_PP-OCRv6_det_medium: + multi_PP-OCRv6_det_medium: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/det/PP-OCRv6_det_medium.onnx SHA256: 92078b7355007ccfffcd4c8cd441a3afd4538904d06881b29a155e1e679907c2 rec: - ch_PP-OCRv6_rec_tiny: + multi_PP-OCRv6_rec_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/rec/PP-OCRv6_rec_tiny.onnx SHA256: e16e242de5937ad92609223f19bc2aff3727ee40b095f996907c24749bad251b - ch_PP-OCRv6_rec_small: + multi_PP-OCRv6_rec_small: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/rec/PP-OCRv6_rec_small.onnx SHA256: 6f327246b50388f3c176ae304bd95767ea6dc0c9ae92153ef8cbe210b3c14884 - ch_PP-OCRv6_rec_medium: + multi_PP-OCRv6_rec_medium: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/rec/PP-OCRv6_rec_medium.onnx SHA256: eef444829dbbe18d7fea59a3f6eb75647518d2b3a9568d27c92e42940204894b openvino: @@ -255,23 +255,23 @@ openvino: SHA256: 7d3c02ef6c7da8ae08b4347cc7695b2081aae68c325d64375724ecf39c99e743 PP-OCRv6: det: - ch_PP-OCRv6_det_tiny: + multi_PP-OCRv6_det_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/det/PP-OCRv6_det_tiny.onnx SHA256: f42c0fbd294d95eac1a550e131b277dac97462c8025fa4b6c3cec1b7894bd3d5 - ch_PP-OCRv6_det_small: + multi_PP-OCRv6_det_small: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/det/PP-OCRv6_det_small.onnx SHA256: 090f04abcd9d9a7498bc4ebf677e4cb9bdce1fe4197ddb7e529f1ef44e1ff94f - ch_PP-OCRv6_det_medium: + multi_PP-OCRv6_det_medium: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/det/PP-OCRv6_det_medium.onnx SHA256: 92078b7355007ccfffcd4c8cd441a3afd4538904d06881b29a155e1e679907c2 rec: - ch_PP-OCRv6_rec_tiny: + multi_PP-OCRv6_rec_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/rec/PP-OCRv6_rec_tiny.onnx SHA256: e16e242de5937ad92609223f19bc2aff3727ee40b095f996907c24749bad251b - ch_PP-OCRv6_rec_small: + multi_PP-OCRv6_rec_small: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/rec/PP-OCRv6_rec_small.onnx SHA256: 6f327246b50388f3c176ae304bd95767ea6dc0c9ae92153ef8cbe210b3c14884 - ch_PP-OCRv6_rec_medium: + multi_PP-OCRv6_rec_medium: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv6/rec/PP-OCRv6_rec_medium.onnx SHA256: eef444829dbbe18d7fea59a3f6eb75647518d2b3a9568d27c92e42940204894b mnn: diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index 75c49f7..6e7aee5 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/python/rapidocr/inference_engine/base.py @@ -11,6 +11,7 @@ import numpy as np from omegaconf import OmegaConf from ..utils.log import logger +from ..utils.model_resolver import normalize_lang, resolve_model_key from ..utils.typings import EngineType, ModelType, OCRVersion, TaskType from ..utils.utils import import_package @@ -77,7 +78,7 @@ class FileInfo: engine_type: EngineType ocr_version: OCRVersion task_type: TaskType - lang_type: Enum + lang_type: Union[Enum, str] model_type: ModelType @@ -117,7 +118,7 @@ class InferSession(abc.ABC): engine_type = file_info.engine_type.value ocr_version = file_info.ocr_version.value task_type = file_info.task_type.value - lang_type = file_info.lang_type.value + lang_type = normalize_lang(file_info.lang_type) model_type = file_info.model_type.value model_dict = OmegaConf.select( @@ -129,6 +130,21 @@ class InferSession(abc.ABC): f"Unsupported configuration: {engine_type}.{ocr_version}.{task_type}.{model_type}" ) + model_key = resolve_model_key( + file_info.task_type, + file_info.ocr_version, + file_info.lang_type, + file_info.model_type, + ) + + if model_key is not None: + if model_key in model_dict: + return model_dict[model_key] + + raise ValueError( + f"Unsupported configuration: {engine_type}.{ocr_version}.{task_type}.{lang_type}.{model_type}" + ) + # 优先查找 server 模型 if model_type == ModelType.SERVER.value: for k in model_dict: diff --git a/python/rapidocr/inference_engine/tensorrt/main.py b/python/rapidocr/inference_engine/tensorrt/main.py index e87e484..1080952 100644 --- a/python/rapidocr/inference_engine/tensorrt/main.py +++ b/python/rapidocr/inference_engine/tensorrt/main.py @@ -11,6 +11,7 @@ from cuda.bindings import runtime as cudart from ...utils.download_file import DownloadFile, DownloadFileInput from ...utils.log import logger +from ...utils.model_resolver import normalize_lang, resolve_model_key from ...utils.typings import EngineType from ..base import FileInfo, InferSession from .engine_builder import TRTEngineBuilder @@ -317,8 +318,17 @@ class TRTInferSession(InferSession): if cfg.get("model_path"): return Path(cfg["model_path"]).stem + model_key = resolve_model_key( + cfg.task_type, + cfg.ocr_version, + cfg.lang_type, + cfg.model_type, + ) + if model_key is not None: + return model_key + task_type = cfg.task_type.value - lang_type = cfg.lang_type.value + lang_type = normalize_lang(cfg.lang_type) ocr_version = cfg.ocr_version.value model_type = cfg.model_type.value return f"{lang_type}_{ocr_version}_{task_type}_{model_type}" diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index 0da2bfe..2291ca6 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -35,6 +35,13 @@ root_dir = Path(__file__).resolve().parent DEFAULT_CFG_PATH = root_dir / "config.yaml" +def normalize_cli_lang_type(lang_type: str) -> Union[LangRec, str]: + try: + return LangRec(lang_type) + except ValueError: + return lang_type + + class RapidOCR: def __init__( self, config_path: Optional[str] = None, params: Optional[Dict[str, Any]] = None @@ -392,7 +399,6 @@ def parse_args(arg_list: Optional[List[str]] = None): "--lang_type", type=str, default="ch", - choices=list(v.value for v in LangRec), ) parser.add_argument("-vis", "--vis_res", action="store_true", default=False) parser.add_argument("--vis_save_dir", type=Path, default=".") @@ -433,10 +439,11 @@ def main(arg_list: Optional[List[str]] = None): if args.command == "check": check_required_files() + lang_type = normalize_cli_lang_type(args.lang_type) params = { "Global.text_score": args.text_score, "Global.return_word_box": args.return_word_box, - "Rec.lang_type": LangRec(args.lang_type), + "Rec.lang_type": lang_type, } ocr_engine = RapidOCR(params=params) @@ -462,7 +469,7 @@ def main(arg_list: Optional[List[str]] = None): vis = VisRes( text_score=args.text_score, font_path=args.font_path, - lang_type=LangRec(args.lang_type), + lang_type=lang_type, ) cur_dir = args.vis_save_dir diff --git a/python/rapidocr/utils/model_resolver.py b/python/rapidocr/utils/model_resolver.py new file mode 100644 index 0000000..3f85cac --- /dev/null +++ b/python/rapidocr/utils/model_resolver.py @@ -0,0 +1,160 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +from dataclasses import dataclass, field +from enum import Enum +from typing import FrozenSet, List, Mapping, Optional, Union + +from .typings import ModelType, OCRVersion, TaskType + + +@dataclass(frozen=True) +class ModelRoute: + model_key_template: str + supported_langs_by_model_type: Mapping[ModelType, FrozenSet[str]] + aliases: Mapping[str, str] = field(default_factory=dict) + + def supported_langs(self, model_type: ModelType) -> FrozenSet[str]: + return self.supported_langs_by_model_type.get(model_type, frozenset()) + + +COMMON_LANG_ALIASES = { + "zh": "ch", + "zh_cn": "ch", + "zh-cn": "ch", + "zh_tw": "chinese_cht", + "zh-tw": "chinese_cht", + "ja": "japan", + "jp": "japan", + "ko": "korean", +} + +PP_OCRV6_LANGS = frozenset( + { + "ch", + "chinese_cht", + "en", + "japan", + "af", + "az", + "bs", + "ca", + "cs", + "cy", + "da", + "de", + "es", + "et", + "eu", + "fi", + "fr", + "ga", + "gl", + "hr", + "hu", + "id", + "is", + "it", + "ku", + "la", + "lb", + "lt", + "lv", + "mi", + "ms", + "mt", + "nl", + "no", + "oc", + "pi", + "pt", + "qu", + "rm", + "ro", + "rs_latin", + "sk", + "sl", + "sq", + "sv", + "sw", + "tl", + "tr", + "uz", + "vi", + "french", + "german", + } +) +PP_OCRV6_TINY_LANGS = PP_OCRV6_LANGS - {"japan"} +PP_OCRV6_LANGS_BY_MODEL_TYPE = { + ModelType.TINY: PP_OCRV6_TINY_LANGS, + ModelType.SMALL: PP_OCRV6_LANGS, + ModelType.MEDIUM: PP_OCRV6_LANGS, +} + +MODEL_ROUTES = { + TaskType.DET: { + OCRVersion.PPOCRV6: ModelRoute( + model_key_template="multi_PP-OCRv6_det_{model_type}", + supported_langs_by_model_type=PP_OCRV6_LANGS_BY_MODEL_TYPE, + aliases=COMMON_LANG_ALIASES, + ), + }, + TaskType.REC: { + OCRVersion.PPOCRV6: ModelRoute( + model_key_template="multi_PP-OCRv6_rec_{model_type}", + supported_langs_by_model_type=PP_OCRV6_LANGS_BY_MODEL_TYPE, + aliases=COMMON_LANG_ALIASES, + ), + }, +} + + +def normalize_lang(lang_type: Union[Enum, str]) -> str: + if isinstance(lang_type, Enum): + lang = lang_type.value + else: + lang = str(lang_type) + + return lang.strip().lower() + + +def resolve_model_key( + task_type: TaskType, + ocr_version: OCRVersion, + lang_type: Union[Enum, str], + model_type: ModelType, +) -> Optional[str]: + route = MODEL_ROUTES.get(task_type, {}).get(ocr_version) + if route is None: + return None + + lang = normalize_lang(lang_type) + lang = route.aliases.get(lang, lang) + supported_langs = route.supported_langs(model_type) + + if lang not in supported_langs: + raise ValueError( + f"Unsupported {task_type.value}.lang_type={lang!r} for " + f"{ocr_version.value} {model_type.value} model." + ) + + return route.model_key_template.format(model_type=model_type.value) + + +def list_supported_langs( + task_type: TaskType, + ocr_version: OCRVersion, + model_type: Optional[ModelType] = None, +) -> List[str]: + route = MODEL_ROUTES.get(task_type, {}).get(ocr_version) + if route is None: + return [] + + if model_type is not None: + return sorted(route.supported_langs(model_type)) + + supported_langs = set() + for langs in route.supported_langs_by_model_type.values(): + supported_langs.update(langs) + return sorted(supported_langs) diff --git a/python/rapidocr/utils/parse_parameters.py b/python/rapidocr/utils/parse_parameters.py index 16d2eb8..8403b7b 100644 --- a/python/rapidocr/utils/parse_parameters.py +++ b/python/rapidocr/utils/parse_parameters.py @@ -39,7 +39,6 @@ class ParseParams(OmegaConf): "engine_type", "model_type", "ocr_version", - "lang_type", "task_type", ] for k, v in params.items(): @@ -75,12 +74,21 @@ class ParseParams(OmegaConf): @staticmethod def LangType(task_type: TaskType, lang_type: str): if task_type == TaskType.DET: - return LangDet(lang_type) + try: + return LangDet(lang_type) + except ValueError: + return lang_type # Return the original string if it's not a valid LangDet value if task_type == TaskType.CLS: - return LangCls(lang_type) + try: + return LangCls(lang_type) + except ValueError: + return lang_type # Return the original string if it's not a valid LangCls value if task_type == TaskType.REC: - return LangRec(lang_type) + try: + return LangRec(lang_type) + except ValueError: + return lang_type # Return the original string if it's not a valid LangRec value raise ValueError(f"task_type {task_type.value} is not in [Det, Cls, Rec]") diff --git a/python/rapidocr/utils/vis_res.py b/python/rapidocr/utils/vis_res.py index 69f4525..215e744 100644 --- a/python/rapidocr/utils/vis_res.py +++ b/python/rapidocr/utils/vis_res.py @@ -3,11 +3,11 @@ # @Contact: liekkaskono@163.com import math import random +from functools import lru_cache from pathlib import Path from typing import List, Optional, Sequence, Tuple, Union import cv2 -from functools import lru_cache import numpy as np from omegaconf import OmegaConf from PIL import Image, ImageDraw, ImageFont @@ -15,6 +15,7 @@ from PIL import Image, ImageDraw, ImageFont from .download_file import DownloadFile, DownloadFileInput from .load_image import LoadImage from .log import logger +from .model_resolver import normalize_lang from .typings import LangRec root_dir = Path(__file__).resolve().parent.parent @@ -32,7 +33,8 @@ def font_cfg(): @lru_cache() def get_font_path( - font_path: Optional[Union[str, Path]] = None, lang_type: Optional[LangRec] = None + font_path: Optional[Union[str, Path]] = None, + lang_type: Union[LangRec, None, str] = None, ) -> str: default_info = font_cfg()["ch"] default_input_params = DownloadFileInput( @@ -47,7 +49,7 @@ def get_font_path( DownloadFile.run(default_input_params) return str(DEFAULT_FONT_PATH) - lang_type = lang_type.value + lang_type = normalize_lang(lang_type) if font_path is None: # 指定了语种，但是没有指定字体文件，根据语种选择字体文件 @@ -60,7 +62,7 @@ def get_font_path( DownloadFile.run(default_input_params) return str(DEFAULT_FONT_PATH) - + font_url, font_sha256 = font_info["path"], font_info["SHA256"] save_font_path = DEFAULT_FONT_DIR / f"{Path(font_url).name}" input_param = DownloadFileInput( @@ -82,7 +84,7 @@ class VisRes: def __init__( self, text_score: float = 0.5, - lang_type: Optional[LangRec] = None, + lang_type: Union[LangRec, None, str] = None, font_path: Optional[str] = None, ): self.text_score = text_score