chore: update files

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

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:

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:

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}"

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
 

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)

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

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