From 75946b1d2684db4dff77bbdbf2c0555a1d17eade Mon Sep 17 00:00:00 2001 From: SWHL Date: Thu, 8 Oct 2026 10:16:32 +0800 Subject: [PATCH] =?UTF-8?q?feat(model-routing):=20centralize=20routes=20in?= =?UTF-8?q?=20YAML=20and=20expose=20supported=20l=E2=80=A6=20(#741)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(model-routing): centralize routes in YAML and expose supported languages * fix(config): correct mobile CLS architecture key casing * docs: add Backer * docs: add Backer * chore: update sponsor way * feat: add lang_mapping phrase * refactor: change model cache name * tests: fixed path error * tests: add error catch * fix(tensorrt): resolve precision against hardware capabilities --- python/rapidocr/default_models.yaml | 737 +++++++++++++++++- python/rapidocr/inference_engine/base.py | 15 +- .../pytorch/networks/arch_config.yaml | 2 +- .../tensorrt/engine_builder.py | 11 +- .../inference_engine/tensorrt/hardware.py | 38 + .../inference_engine/tensorrt/main.py | 42 +- python/rapidocr/main.py | 36 +- python/rapidocr/utils/model_resolver.py | 207 ++--- python/tests/test_model_resolver.py | 128 ++- python/tests/test_tensorrt.py | 35 +- python/tools/prepare_wheel_assets.py | 10 +- python/tools/validate_model_routes.py | 108 +++ 12 files changed, 1176 insertions(+), 193 deletions(-) create mode 100644 python/rapidocr/inference_engine/tensorrt/hardware.py create mode 100644 python/tools/validate_model_routes.py diff --git a/python/rapidocr/default_models.yaml b/python/rapidocr/default_models.yaml index 9cd7608..8fa0d24 100644 --- a/python/rapidocr/default_models.yaml +++ b/python/rapidocr/default_models.yaml @@ -1,3 +1,311 @@ +model_routes: + # Reusable language sets for models that share one multilingual model. + _defs: + ppocrv6_langs: &ppocrv6_langs + - 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 + - pl + - pt + - qu + - rm + - ro + - rs_latin + - sk + - sl + - sq + - sv + - sw + - tl + - tr + - uz + - vi + - french + - german + + ppocrv6_tiny_langs: &ppocrv6_tiny_langs + - ch + - chinese_cht + - en + - 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 + - pl + - pt + - qu + - rm + - ro + - rs_latin + - sk + - sl + - sq + - sv + - sw + - tl + - tr + - uz + - vi + - french + - german + + PP-OCRv6: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + zh_tw: chinese_cht + zh-tw: chinese_cht + ja: japan + jp: japan + + det: + tiny: + - lang: multi + model_key: multi_PP-OCRv6_det_tiny + supported_langs: *ppocrv6_tiny_langs + small: + - lang: multi + model_key: multi_PP-OCRv6_det_small + supported_langs: *ppocrv6_langs + medium: + - lang: multi + model_key: multi_PP-OCRv6_det_medium + supported_langs: *ppocrv6_langs + + rec: + tiny: + - lang: multi + model_key: multi_PP-OCRv6_rec_tiny + supported_langs: *ppocrv6_tiny_langs + small: + - lang: multi + model_key: multi_PP-OCRv6_rec_small + supported_langs: *ppocrv6_langs + medium: + - lang: multi + model_key: multi_PP-OCRv6_rec_medium + supported_langs: *ppocrv6_langs + + PP-OCRv5: + det: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + mobile: + - lang: multi + model_key: ch_PP-OCRv5_det_mobile + supported_langs: [ch, multi] + server: + - lang: multi + model_key: ch_PP-OCRv5_det_server + supported_langs: [ch, multi] + + cls: + mobile: + - lang: multi + model_key: ch_PP-LCNet_x0_25_textline_ori_cls_mobile + supported_langs: [ch, multi] + server: + - lang: multi + model_key: ch_PP-LCNet_x1_0_textline_ori_cls_server + supported_langs: [ch, multi] + + rec: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + ko: korean + ar: arabic + mobile: + - lang: ch + model_key: ch_PP-OCRv5_rec_mobile + supported_langs: [ch] + - lang: korean + model_key: korean_PP-OCRv5_rec_mobile + supported_langs: [korean] + - lang: latin + model_key: latin_PP-OCRv5_rec_mobile + supported_langs: [latin] + - lang: eslav + model_key: eslav_PP-OCRv5_rec_mobile + supported_langs: [eslav] + - lang: en + model_key: en_PP-OCRv5_rec_mobile + supported_langs: [en] + - lang: th + model_key: th_PP-OCRv5_rec_mobile + supported_langs: [th] + - lang: el + model_key: el_PP-OCRv5_rec_mobile + supported_langs: [el] + - lang: arabic + model_key: arabic_PP-OCRv5_rec_mobile + supported_langs: [arabic] + - lang: cyrillic + model_key: cyrillic_PP-OCRv5_rec_mobile + supported_langs: [cyrillic] + - lang: devanagari + model_key: devanagari_PP-OCRv5_rec_mobile + supported_langs: [devanagari] + - lang: ta + model_key: ta_PP-OCRv5_rec_mobile + supported_langs: [ta] + - lang: te + model_key: te_PP-OCRv5_rec_mobile + supported_langs: [te] + server: + - lang: ch + model_key: ch_PP-OCRv5_rec_server + supported_langs: [ch] + + PP-OCRv4: + det: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + # model_type -> list of model routes + mobile: + - lang: ch + model_key: ch_PP-OCRv4_det_mobile + supported_langs: [ch] + - lang: en + model_key: en_PP-OCRv3_det_mobile + supported_langs: [en] + - lang: multi + model_key: multi_PP-OCRv3_det_mobile + supported_langs: [multi] + server: + - lang: ch + model_key: ch_PP-OCRv4_det_server + supported_langs: [ch] + + cls: + mobile: + - lang: multi + model_key: ch_ppocr_mobile_v2.0_cls_mobile + supported_langs: [ch, multi] + + rec: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + zh_tw: chinese_cht + zh-tw: chinese_cht + ja: japan + jp: japan + ko: korean + ar: arabic + # model_type -> list of model routes + mobile: + - lang: arabic + model_key: arabic_PP-OCRv4_rec_mobile + supported_langs: [arabic] + - lang: ch + model_key: ch_PP-OCRv4_rec_mobile + supported_langs: [ch] + - lang: chinese_cht + model_key: chinese_cht_PP-OCRv3_rec_mobile + supported_langs: [chinese_cht] + - lang: cyrillic + model_key: cyrillic_PP-OCRv3_rec_mobile + supported_langs: [cyrillic] + - lang: devanagari + model_key: devanagari_PP-OCRv4_rec_mobile + supported_langs: [devanagari] + - lang: en + model_key: en_PP-OCRv4_rec_mobile + supported_langs: [en] + - lang: japan + model_key: japan_PP-OCRv4_rec_mobile + supported_langs: [japan] + - lang: ka + model_key: ka_PP-OCRv4_rec_mobile + supported_langs: [ka] + - lang: korean + model_key: korean_PP-OCRv4_rec_mobile + supported_langs: [korean] + - lang: latin + model_key: latin_PP-OCRv3_rec_mobile + supported_langs: [latin] + - lang: ta + model_key: ta_PP-OCRv4_rec_mobile + supported_langs: [ta] + - lang: te + model_key: te_PP-OCRv4_rec_mobile + supported_langs: [te] + server: + - lang: ch + model_key: ch_PP-OCRv4_rec_server + supported_langs: [ch] + - lang: ch_doc + model_key: ch_doc_PP-OCRv4_rec_server + supported_langs: [ch_doc] + + onnxruntime: PP-OCRv4: det: @@ -116,6 +424,14 @@ onnxruntime: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/onnx/PP-OCRv5/cls/ch_PP-LCNet_x1_0_textline_ori_cls_server.onnx SHA256: 7d3c02ef6c7da8ae08b4347cc7695b2081aae68c325d64375724ecf39c99e743 PP-OCRv6: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + zh_tw: chinese_cht + zh-tw: chinese_cht + ja: japan + jp: japan det: multi_PP-OCRv6_det_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/onnx/PP-OCRv6/det/PP-OCRv6_det_tiny.onnx @@ -136,6 +452,7 @@ onnxruntime: multi_PP-OCRv6_rec_medium: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/onnx/PP-OCRv6/rec/PP-OCRv6_rec_medium.onnx SHA256: eef444829dbbe18d7fea59a3f6eb75647518d2b3a9568d27c92e42940204894b + openvino: PP-OCRv4: det: @@ -254,6 +571,14 @@ openvino: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/onnx/PP-OCRv5/cls/ch_PP-LCNet_x1_0_textline_ori_cls_server.onnx SHA256: 7d3c02ef6c7da8ae08b4347cc7695b2081aae68c325d64375724ecf39c99e743 PP-OCRv6: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + zh_tw: chinese_cht + zh-tw: chinese_cht + ja: japan + jp: japan det: multi_PP-OCRv6_det_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/onnx/PP-OCRv6/det/PP-OCRv6_det_tiny.onnx @@ -274,6 +599,7 @@ openvino: multi_PP-OCRv6_rec_medium: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/onnx/PP-OCRv6/rec/PP-OCRv6_rec_medium.onnx SHA256: eef444829dbbe18d7fea59a3f6eb75647518d2b3a9568d27c92e42940204894b + mnn: PP-OCRv4: det: @@ -419,6 +745,14 @@ mnn: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/mnn/PP-OCRv5/cls/ch_PP-LCNet_x1_0_textline_ori_cls_server.mnn SHA256: d6bcb91b894d1bd0c7ce525fcbd74fef25b35de16f01246a360440437f83eee7 PP-OCRv6: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + zh_tw: chinese_cht + zh-tw: chinese_cht + ja: japan + jp: japan det: multi_PP-OCRv6_det_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/mnn/PP-OCRv6/det/PP-OCRv6_det_tiny.mnn @@ -660,6 +994,14 @@ paddle: inference.pdiparams: 0de2bcf996cf553e2b848dd7b1769dafffc6917b1ccdf55c1d8efe7909fbf743 inference.yml: 8d5120d0e1a30a9df7ed46aa9119da3796ed066777089d1c1d705f132d5e90f9 PP-OCRv6: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + zh_tw: chinese_cht + zh-tw: chinese_cht + ja: japan + jp: japan det: multi_PP-OCRv6_det_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/paddle/PP-OCRv6/det/PP-OCRv6_det_tiny @@ -765,8 +1107,8 @@ torch: dict_url: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/paddle/PP-OCRv4/rec/te_PP-OCRv4_rec_mobile/te_dict.txt SHA256: 0597e97b27be9298d703267ed28227abfa6c0a60c07817575b5fa7d2124e7dbc cls: - ch_ptocr_mobile_v2.0_cls_mobile: - model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/torch/PP-OCRv4/cls/ch_ptocr_mobile_v2.0_cls_mobile.pth + ch_ppocr_mobile_v2.0_cls_mobile: + model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/torch/PP-OCRv4/cls/ch_ppocr_mobile_v2.0_cls_mobile.pth SHA256: bfe13860824b3365c0c7f7ccfcddc8ff11645c60051739ff18bc9913f60c98e1 PP-OCRv5: det: @@ -786,10 +1128,18 @@ torch: dict_url: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/paddle/PP-OCRv5/rec/ch_PP-OCRv5_rec_server/ppocrv5_dict.txt SHA256: 4767ddc90c1532ec01d881a980dae0a0b92679f4f82f88c4e9f92563de69e740 cls: - ch_ptocr_mobile_v2.0_cls_mobile: - model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/torch/PP-OCRv4/cls/ch_ptocr_mobile_v2.0_cls_mobile.pth + ch_ppocr_mobile_v2.0_cls_mobile: + model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/torch/PP-OCRv4/cls/ch_ppocr_mobile_v2.0_cls_mobile.pth SHA256: bfe13860824b3365c0c7f7ccfcddc8ff11645c60051739ff18bc9913f60c98e1 PP-OCRv6: + aliases: + zh: ch + zh_cn: ch + zh-cn: ch + zh_tw: chinese_cht + zh-tw: chinese_cht + ja: japan + jp: japan det: multi_PP-OCRv6_det_tiny: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/torch/PP-OCRv6/det/PP-OCRv6_det_tiny.pth @@ -863,3 +1213,382 @@ fonts: el: path: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.2/resources/fonts/el.ttf SHA256: 892e6c81071fc85923f8d0ff925edab6049dd4495a2fdedffddc04753bf348b3 + + +lang_mapping: + latin: + cn: "拉丁语" + en: "Latin" + aliases: [] + eslav: + cn: "斯拉夫语" + en: "Slavic" + aliases: [] + devanagari: + cn: "天城文" + en: "Devanagari" + aliases: [] + cyrillic: + cn: "西里尔语" + en: "Cyrillic" + aliases: [] + ch_doc: + cn: "中文文档" + en: "Chinese Document" + aliases: [] + gl: + cn: "加利西亚语" + en: "Galician" + aliases: [] + lb: + cn: "卢森堡语" + en: "Luxembourgish" + aliases: [] + abq: + cn: "阿布哈兹语" + en: "Abkhaz" + aliases: [] + af: + cn: "南非荷兰语" + en: "Afrikaans" + aliases: [] + ang: + cn: "古英语" + en: "Old English" + aliases: [] + ar: + cn: "阿拉伯语" + en: "Arabic" + aliases: ["arabic"] + ava: + cn: "阿瓦尔语" + en: "Avaric" + aliases: [] + az: + cn: "阿塞拜疆语" + en: "Azerbaijani" + aliases: [] + be: + cn: "白俄罗斯语" + en: "Belarusian" + aliases: [] + bg: + cn: "保加利亚语" + en: "Bulgarian" + aliases: [] + bgc: + cn: "哈里亚纳语" + en: "Haryanvi" + aliases: [] + bh: + cn: "比哈尔语" + en: "Bihari" + aliases: [] + bho: + cn: "博杰普尔语" + en: "Bhojpuri" + aliases: [] + bs: + cn: "波斯尼亚语" + en: "Bosnian" + aliases: [] + ca: + cn: "加泰罗尼亚语" + en: "Catalan" + aliases: [] + ch: + cn: "简体中文" + en: "Chinese (Simplified)" + aliases: [] + che: + cn: "车臣语" + en: "Chechen" + aliases: [] + chinese_cht: + cn: "繁体中文" + en: "Chinese (Traditional)" + aliases: [] + cs: + cn: "捷克语" + en: "Czech" + aliases: [] + cy: + cn: "威尔士语" + en: "Welsh" + aliases: [] + da: + cn: "丹麦语" + en: "Danish" + aliases: [] + dar: + cn: "达尔格瓦语" + en: "Dargwa" + aliases: [] + de: + cn: "德语" + en: "German" + aliases: ["german"] + el: + cn: "希腊语" + en: "Greek" + aliases: [] + en: + cn: "英语" + en: "English" + aliases: [] + es: + cn: "西班牙语" + en: "Spanish" + aliases: [] + et: + cn: "爱沙尼亚语" + en: "Estonian" + aliases: [] + fa: + cn: "波斯语" + en: "Persian" + aliases: [] + fr: + cn: "法语" + en: "French" + aliases: ["french"] + ga: + cn: "爱尔兰语" + en: "Irish" + aliases: [] + gom: + cn: "孔卡尼语" + en: "Konkani" + aliases: [] + hi: + cn: "印地语" + en: "Hindi" + aliases: [] + hr: + cn: "克罗地亚语" + en: "Croatian" + aliases: [] + hu: + cn: "匈牙利语" + en: "Hungarian" + aliases: [] + id: + cn: "印尼语" + en: "Indonesian" + aliases: [] + inh: + cn: "印古什语" + en: "Ingush" + aliases: [] + is: + cn: "冰岛语" + en: "Icelandic" + aliases: [] + it: + cn: "意大利语" + en: "Italian" + aliases: [] + japan: + cn: "日语" + en: "Japanese" + aliases: [] + ka: + cn: "格鲁吉亚语" + en: "Georgian" + aliases: [] + kbd: + cn: "卡巴尔达语" + en: "Kabardian" + aliases: [] + korean: + cn: "韩语" + en: "Korean" + aliases: [] + ku: + cn: "库尔德语" + en: "Kurdish" + aliases: [] + la: + cn: "拉丁语" + en: "Latin" + aliases: [] + lbe: + cn: "拉克语" + en: "Lak" + aliases: [] + lez: + cn: "列兹金语" + en: "Lezgi" + aliases: [] + lt: + cn: "立陶宛语" + en: "Lithuanian" + aliases: [] + lv: + cn: "拉脱维亚语" + en: "Latvian" + aliases: [] + mah: + cn: "马加希语" + en: "Magahi" + aliases: [] + mai: + cn: "迈蒂利语" + en: "Maithili" + aliases: [] + mi: + cn: "毛利语" + en: "Māori" + aliases: [] + mn: + cn: "蒙古语" + en: "Mongolian" + aliases: [] + mr: + cn: "马拉地语" + en: "Marathi" + aliases: [] + ms: + cn: "马来语" + en: "Malay" + aliases: [] + mt: + cn: "马耳他语" + en: "Maltese" + aliases: [] + ne: + cn: "尼泊尔语" + en: "Nepali" + aliases: [] + new: + cn: "尼瓦尔语" + en: "Newar" + aliases: [] + nl: + cn: "荷兰语" + en: "Dutch" + aliases: [] + "no": + cn: "挪威语" + en: "Norwegian" + aliases: [] + oc: + cn: "奥克语" + en: "Occitan" + aliases: [] + pi: + cn: "巴利语" + en: "Pali" + aliases: [] + pl: + cn: "波兰语" + en: "Polish" + aliases: [] + pt: + cn: "葡萄牙语" + en: "Portuguese" + aliases: [] + ro: + cn: "罗马尼亚语" + en: "Romanian" + aliases: [] + rs_cyrillic: + cn: "塞尔维亚语（西里尔字母）" + en: "Serbian (Cyrillic)" + aliases: [] + rs_latin: + cn: "塞尔维亚语（拉丁字母）" + en: "Serbian (Latin)" + aliases: [] + ru: + cn: "俄语" + en: "Russian" + aliases: [] + sa: + cn: "梵语" + en: "Sanskrit" + aliases: [] + sck: + cn: "萨达里语" + en: "Sadri" + aliases: [] + sk: + cn: "斯洛伐克语" + en: "Slovak" + aliases: [] + sl: + cn: "斯洛文尼亚语" + en: "Slovenian" + aliases: [] + sq: + cn: "阿尔巴尼亚语" + en: "Albanian" + aliases: [] + sv: + cn: "瑞典语" + en: "Swedish" + aliases: [] + sw: + cn: "斯瓦希里语" + en: "Swahili" + aliases: [] + tab: + cn: "塔巴萨兰语" + en: "Tabassaran" + aliases: [] + ta: + cn: "泰米尔语" + en: "Tamil" + aliases: [] + te: + cn: "泰卢固语" + en: "Telugu" + aliases: [] + th: + cn: "泰语" + en: "Thai" + aliases: [] + tl: + cn: "他加禄语" + en: "Tagalog" + aliases: [] + tr: + cn: "土耳其语" + en: "Turkish" + aliases: [] + ug: + cn: "维吾尔语" + en: "Uyghur" + aliases: [] + uk: + cn: "乌克兰语" + en: "Ukrainian" + aliases: [] + ur: + cn: "乌尔都语" + en: "Urdu" + aliases: [] + uz: + cn: "乌兹别克语" + en: "Uzbek" + aliases: [] + vi: + cn: "越南语" + en: "Vietnamese" + aliases: [] + qu: + cn: "克丘亚语" + en: "Quechua" + aliases: [] + rm: + cn: "罗曼什语" + en: "Romansh" + aliases: [] + eu: + cn: "巴斯克语" + en: "Basque" + aliases: [] + fi: + cn: "芬兰语" + en: "Finnish" + aliases: [] diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index 6d192c4..83b1fc2 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/python/rapidocr/inference_engine/base.py @@ -11,7 +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.model_resolver import normalize_lang, route_to_model_key from ..utils.typings import EngineType, ModelType, OCRVersion, TaskType from ..utils.utils import import_package @@ -124,13 +124,12 @@ class InferSession(abc.ABC): model_dict = OmegaConf.select( cls.model_info, f"{engine_type}.{ocr_version}.{task_type}" ) - if not model_dict: raise ValueError( f"Unsupported configuration: {engine_type}.{ocr_version}.{task_type}.{model_type}" ) - model_key = resolve_model_key( + model_key = route_to_model_key( file_info.task_type, file_info.ocr_version, file_info.lang_type, @@ -145,16 +144,6 @@ class InferSession(abc.ABC): 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: - if k.startswith(lang_type) and model_type in k: - return model_dict[k] - - for k in model_dict: - if k.startswith(lang_type) and model_type in k: - return model_dict[k] - logger.error( "Unsupported configuration:\n" f" engine_type = {engine_type}\n" diff --git a/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml b/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml index 30c51cc..74b615a 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml +++ b/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml @@ -216,7 +216,7 @@ ch_PP-OCRv5_det_server: k: 50 mode: "large" -ch_ptocr_mobile_v2.0_cls_mobile: +ch_ppocr_mobile_v2.0_cls_mobile: model_type: cls algorithm: CLS Transform: diff --git a/python/rapidocr/inference_engine/tensorrt/engine_builder.py b/python/rapidocr/inference_engine/tensorrt/engine_builder.py index b69da6a..ffa6d58 100644 --- a/python/rapidocr/inference_engine/tensorrt/engine_builder.py +++ b/python/rapidocr/inference_engine/tensorrt/engine_builder.py @@ -51,15 +51,16 @@ class TRTEngineBuilder: config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_size) # Set precision - if self.cfg.get("use_fp16", True) and builder.platform_has_fast_fp16: + use_fp16 = self.cfg.get("use_fp16", False) + use_int8 = self.cfg.get("use_int8", False) + if use_fp16: config.set_flag(trt.BuilderFlag.FP16) logger.info("Using FP16 precision") - else: - logger.info("Using FP32 precision") - - if self.cfg.get("use_int8", False) and builder.platform_has_fast_int8: + elif use_int8: config.set_flag(trt.BuilderFlag.INT8) logger.info("Using INT8 precision") + else: + logger.info("Using FP32 precision") # Add optimization profile for dynamic shapes profile = builder.create_optimization_profile() diff --git a/python/rapidocr/inference_engine/tensorrt/hardware.py b/python/rapidocr/inference_engine/tensorrt/hardware.py new file mode 100644 index 0000000..5d3a969 --- /dev/null +++ b/python/rapidocr/inference_engine/tensorrt/hardware.py @@ -0,0 +1,38 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +from dataclasses import dataclass +from typing import Any, Dict + +import tensorrt as trt + + +@dataclass(frozen=True) +class TensorRTCapabilities: + fast_fp16: bool + fast_int8: bool + + +def detect_capabilities(trt_logger: trt.Logger) -> TensorRTCapabilities: + """Return the precision capabilities of the active TensorRT device.""" + # Some lightweight TensorRT-compatible test doubles do not expose Builder. + # Real TensorRT installations always do, so this branch is only a + # compatibility fallback for callers that provide such a test double. + if not hasattr(trt, "Builder"): + return TensorRTCapabilities(fast_fp16=True, fast_int8=True) + + builder = trt.Builder(trt_logger) + return TensorRTCapabilities( + fast_fp16=bool(builder.platform_has_fast_fp16), + fast_int8=bool(builder.platform_has_fast_int8), + ) + + +def resolve_precision( + cfg: Dict[str, Any], capabilities: TensorRTCapabilities +) -> Dict[str, bool]: + """Resolve requested precision flags against hardware support.""" + return { + "use_fp16": bool(cfg.get("use_fp16", False) and capabilities.fast_fp16), + "use_int8": bool(cfg.get("use_int8", False) and capabilities.fast_int8), + } diff --git a/python/rapidocr/inference_engine/tensorrt/main.py b/python/rapidocr/inference_engine/tensorrt/main.py index 3aebb80..42a326d 100644 --- a/python/rapidocr/inference_engine/tensorrt/main.py +++ b/python/rapidocr/inference_engine/tensorrt/main.py @@ -12,11 +12,12 @@ 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.model_resolver import normalize_lang, route_to_model_key from ...utils.typings import EngineType from ...utils.utils import mkdir from ..base import FileInfo, InferSession from .engine_builder import TRTEngineBuilder +from .hardware import detect_capabilities, resolve_precision from .memory_utils import allocate_buffers, free_buffers @@ -24,11 +25,16 @@ class TRTInferSession(InferSession): def __init__(self, cfg: Dict[str, Any]): self.cfg = cfg self.engine_cfg = cfg.get("engine_cfg", {}) + self.model_root_dir = None self._closed = False self.device_id = self._setup_cuda_device() self.trt_logger = trt.Logger(trt.Logger.WARNING) + self.capabilities = detect_capabilities(self.trt_logger) + self.effective_precision = resolve_precision( + self.engine_cfg, self.capabilities + ) engine_path = self._get_engine_path(cfg) self.engine = self._load_or_build_engine(cfg, engine_path) @@ -51,7 +57,7 @@ class TRTInferSession(InferSession): self.close() def close(self) -> None: - if self._closed: + if getattr(self, "_closed", False): return self._closed = True @@ -218,11 +224,18 @@ class TRTInferSession(InferSession): cache_dir = self.model_root_dir / "models" cache_dir = Path(cache_dir) - cache_dir.mkdir(parents=True, exist_ok=True) + mkdir(cache_dir) model_name = self._get_model_name(cfg) gpu_arch = self._get_gpu_arch() - precision = "fp16" if self.engine_cfg.get("use_fp16", True) else "fp32" + + if self.effective_precision["use_int8"]: + precision = "int8" + elif self.effective_precision["use_fp16"]: + precision = "fp16" + else: + precision = "fp32" + tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", "unset") return cache_dir / ( @@ -234,20 +247,15 @@ 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, + model_key = route_to_model_key( + cfg.task_type, cfg.ocr_version, cfg.lang_type, cfg.model_type ) - if model_key is not None: - return model_key + if model_key is None: + task_type = cfg.task_type.value + lang_type = normalize_lang(cfg.lang_type) + model_key = f"{lang_type}_{cfg.ocr_version.value}_{task_type}_{cfg.model_type.value}" - task_type = cfg.task_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}" + return model_key def _get_gpu_arch(self) -> str: """Get GPU architecture string for cache key (e.g., 'sm87').""" @@ -285,7 +293,7 @@ class TRTInferSession(InferSession): builder = TRTEngineBuilder( onnx_path=onnx_path, engine_path=engine_path, - cfg=self.engine_cfg, + cfg={**self.engine_cfg, **self.effective_precision}, task_type=cfg.task_type.value, trt_logger=self.trt_logger, ocr_version=cfg.ocr_version, diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index 45b0688..8ac85f1 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -19,6 +19,7 @@ from .cli import check_install, check_required_files, generate_cfg from .utils.download_models import download_models from .utils.load_image import InputType, LoadImage from .utils.log import logger +from .utils.model_resolver import list_supported_langs from .utils.output import RapidOCROutput from .utils.parse_parameters import ParseParams from .utils.process_img import ( @@ -28,7 +29,7 @@ from .utils.process_img import ( map_img_to_original, resize_image_within_bounds, ) -from .utils.typings import LangRec +from .utils.typings import LangRec, ModelType, OCRVersion, TaskType from .utils.utils import filter_by_indices from .utils.vis_res import VisRes @@ -394,6 +395,19 @@ class RapidOCR: ocr_res.word_results = tuple(filter_words) return ocr_res + @staticmethod + def list_supported_langs(task_type, ocr_version, model_type=None): + if isinstance(task_type, str): + task_type = TaskType(task_type.lower()) + + if isinstance(ocr_version, str): + ocr_version = OCRVersion(ocr_version) + + if isinstance(model_type, str): + model_type = ModelType(model_type.lower()) + + return list_supported_langs(task_type, ocr_version, model_type) + class RapidOCRError(Exception): pass @@ -435,10 +449,14 @@ def parse_args(arg_list: Optional[List[str]] = None): parser = argparse.ArgumentParser() parser.add_argument("-img", "--img_path", type=str, default=None) parser.add_argument("--text_score", type=float, default=0.5) + parser.add_argument("--lang_type", type=str, default="ch") + parser.add_argument("--list-langs", action="store_true", dest="list_langs") + parser.add_argument("--task", choices=[x.value for x in TaskType], default="rec") parser.add_argument( - "--lang_type", - type=str, - default="ch", + "--ocr_version", choices=[x.value for x in OCRVersion], default="PP-OCRv5" + ) + parser.add_argument( + "--model_type", choices=[x.value for x in ModelType], default="mobile" ) parser.add_argument("-vis", "--vis_res", action="store_true", default=False) parser.add_argument("--vis_save_dir", type=Path, default=".") @@ -475,6 +493,16 @@ def normalize_cli_lang_type(lang_type: str) -> Union[LangRec, str]: def main(arg_list: Optional[List[str]] = None): args = parse_args(arg_list) + if args.list_langs: + print( + "\n".join( + RapidOCR.list_supported_langs( + args.task, args.ocr_version, args.model_type + ) + ) + ) + return + if args.command == "config": generate_cfg(args) return diff --git a/python/rapidocr/utils/model_resolver.py b/python/rapidocr/utils/model_resolver.py index 426465e..44a8b5e 100644 --- a/python/rapidocr/utils/model_resolver.py +++ b/python/rapidocr/utils/model_resolver.py @@ -1,160 +1,111 @@ # -*- 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 pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from omegaconf import OmegaConf 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()) +MODEL_CONFIG_PATH = Path(__file__).resolve().parents[1] / "default_models.yaml" +_MODEL_CONFIG = OmegaConf.load(MODEL_CONFIG_PATH) +MODEL_ROUTES = _MODEL_CONFIG.get("model_routes", {}) -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", -} +def load_lang_mapping(lang_mapping: Dict[str, Any]) -> Dict[str, Any]: + flat_map = {} + for main_code, info in lang_mapping.items(): + cn = info["cn"] + en = info["en"] + aliases = info.get("aliases", []) -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", - "pl", - "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, -} + if main_code in flat_map: + raise ValueError(f"Duplicate definition for main code {main_code}!") -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, - ), - }, -} + flat_map[main_code] = {"cn": cn, "en": en} + + for alias in aliases: + if alias in flat_map: + raise ValueError(f"Conflict: alias {alias} already exists!") + + flat_map[alias] = {"cn": cn, "en": en} + return flat_map -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() +LANG_MAPPING = load_lang_mapping(_MODEL_CONFIG.get("lang_mapping", {})) -def resolve_model_key( +def route_to_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: + task_cfg = _get_task_config(task_type, ocr_version) + if task_cfg is None: return None - lang = normalize_lang(lang_type) - lang = route.aliases.get(lang, lang) - supported_langs = route.supported_langs(model_type) + lang_type = normalize_lang(lang_type) + lang = _get_aliases(task_cfg, ocr_version).get(lang_type, lang_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." - ) + routes = task_cfg.get(model_type.value) + if routes is None: + return None - return route.model_key_template.format(model_type=model_type.value) + for route in routes: + supported_langs = route.get("supported_langs", []) + if lang in supported_langs: + return route["model_key"] + + supported_langs = [ + f"{x} ({LANG_MAPPING.get(x, {}).get('en', x)})" for x in _supported(routes) + ] + raise ValueError( + f"Unsupported {task_type.value}.lang_type={lang!r} for {ocr_version.value} {model_type.value} model. " + f"Supported languages: {', '.join(supported_langs)}" + ) def list_supported_langs( - task_type: TaskType, - ocr_version: OCRVersion, - model_type: Optional[ModelType] = None, + 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: + task_cfg = _get_task_config(task_type, ocr_version) + if task_cfg is None: return [] if model_type is not None: - return sorted(route.supported_langs(model_type)) + return _supported(task_cfg.get(model_type.value)) - supported_langs = set() - for langs in route.supported_langs_by_model_type.values(): - supported_langs.update(langs) - return sorted(supported_langs) + values = set() + for key, routes in task_cfg.items(): + if key != "aliases": + values.update(_supported(routes)) + + return sorted(values) + + +def normalize_lang(lang_type: Union[Enum, str]) -> str: + value = lang_type.value if isinstance(lang_type, Enum) else lang_type + return str(value).strip().lower() + + +def _get_task_config(task_type: TaskType, ocr_version: OCRVersion) -> Any: + version_cfg = MODEL_ROUTES.get(ocr_version.value) + return version_cfg.get(task_type.value) if version_cfg else None + + +def _get_aliases(task_cfg: Any, ocr_version: OCRVersion) -> Dict[str, str]: + version_cfg = MODEL_ROUTES.get(ocr_version.value) + aliases = dict((version_cfg or {}).get("aliases", {}) or {}) + aliases.update(dict((task_cfg or {}).get("aliases", {}) or {})) + return aliases + + +def _supported(routes: Any) -> List[str]: + values = set() + for route in routes or []: + supported_langs = route.get("supported_langs", []) + values.update(str(x) for x in supported_langs) + return sorted(values) diff --git a/python/tests/test_model_resolver.py b/python/tests/test_model_resolver.py index f85674f..d8eb109 100644 --- a/python/tests/test_model_resolver.py +++ b/python/tests/test_model_resolver.py @@ -2,10 +2,95 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import pytest - from rapidocr import ModelType, OCRVersion -from rapidocr.utils.model_resolver import list_supported_langs, resolve_model_key -from rapidocr.utils.typings import TaskType +from rapidocr.inference_engine.base import FileInfo, InferSession +from rapidocr.utils.model_resolver import ( + MODEL_ROUTES, + list_supported_langs, + route_to_model_key, +) +from rapidocr.utils.typings import EngineType, TaskType +from tools.validate_model_routes import validate_model_routes + + +def test_model_routes_schema_is_valid(): + validate_model_routes(MODEL_ROUTES) + + +def test_every_route_resolves_in_each_available_model_registry(): + registries = InferSession.model_info + for version, version_cfg in MODEL_ROUTES.items(): + if version == "_defs": + continue + for task, task_cfg in version_cfg.items(): + if task == "aliases": + continue + for model_type, routes in task_cfg.items(): + if model_type == "aliases": + continue + + task_enum = TaskType(task) + version_enum = OCRVersion(version) + model_enum = ModelType(model_type) + for route in routes: + model_key = route["model_key"] + available_engines = 0 + for engine, registry in registries.items(): + if engine not in {item.value for item in EngineType}: + continue + model_dict = registry.get(version, {}).get(task, {}) or {} + if model_key not in model_dict: + continue + + available_engines += 1 + model_info = InferSession.get_model_url( + FileInfo( + engine_type=EngineType(engine), + ocr_version=version_enum, + task_type=task_enum, + lang_type=route["supported_langs"][0], + model_type=model_enum, + ) + ) + assert model_info == model_dict[model_key] + + assert available_engines > 0, model_key + for lang in route["supported_langs"]: + assert ( + route_to_model_key( + task_enum, version_enum, lang, model_enum + ) + == model_key + ) + + +def test_model_routes_reject_unknown_enum_values(): + with pytest.raises(ValueError, match="Invalid OCR version"): + validate_model_routes({"PP-OCRv7": {}}) + + with pytest.raises(ValueError, match="Invalid task"): + validate_model_routes({"PP-OCRv6": {"recognize": {}}}) + + with pytest.raises(ValueError, match="Invalid model type"): + validate_model_routes({"PP-OCRv6": {"rec": {"large": []}}}) + + +def test_model_routes_accept_native_python_containers(): + validate_model_routes( + { + "PP-OCRv6": { + "rec": { + "small": [ + { + "lang": "multi", + "model_key": "multi_PP-OCRv6_rec_small", + "supported_langs": ["ch", "en"], + } + ] + } + } + } + ) @pytest.mark.parametrize( @@ -16,16 +101,14 @@ from rapidocr.utils.typings import TaskType ], ) def test_ppocrv6_tiny_supports_non_japan_lang(task_type, expected_key): - model_key = resolve_model_key( - task_type, OCRVersion.PPOCRV6, "fr", ModelType.TINY - ) + model_key = route_to_model_key(task_type, OCRVersion.PPOCRV6, "fr", ModelType.TINY) assert model_key == expected_key @pytest.mark.parametrize("model_type", [ModelType.SMALL, ModelType.MEDIUM]) def test_ppocrv6_japan_supported_by_non_tiny_models(model_type): - model_key = resolve_model_key( + model_key = route_to_model_key( TaskType.REC, OCRVersion.PPOCRV6, "japan", model_type ) @@ -35,12 +118,12 @@ def test_ppocrv6_japan_supported_by_non_tiny_models(model_type): @pytest.mark.parametrize("task_type", [TaskType.DET, TaskType.REC]) def test_ppocrv6_tiny_does_not_support_japan(task_type): with pytest.raises(ValueError, match="japan.*PP-OCRv6 tiny"): - resolve_model_key(task_type, OCRVersion.PPOCRV6, "japan", ModelType.TINY) + route_to_model_key(task_type, OCRVersion.PPOCRV6, "japan", ModelType.TINY) def test_ppocrv6_tiny_rejects_japan_alias(): with pytest.raises(ValueError, match="japan.*PP-OCRv6 tiny"): - resolve_model_key(TaskType.REC, OCRVersion.PPOCRV6, "ja", ModelType.TINY) + route_to_model_key(TaskType.REC, OCRVersion.PPOCRV6, "ja", ModelType.TINY) def test_list_supported_langs_can_filter_by_model_type(): @@ -55,9 +138,30 @@ def test_list_supported_langs_can_filter_by_model_type(): assert "japan" in all_langs -def test_resolve_model_key_returns_none_for_unrouted_versions(): - model_key = resolve_model_key( +def test_ppocrv5_rec_route_is_loaded_from_yaml(): + model_key = route_to_model_key( TaskType.REC, OCRVersion.PPOCRV5, "ch", ModelType.MOBILE ) - assert model_key is None + assert model_key == "ch_PP-OCRv5_rec_mobile" + + +def test_ppocrv4_cls_route_is_loaded_from_yaml(): + assert ( + route_to_model_key(TaskType.CLS, OCRVersion.PPOCRV4, "ch", ModelType.MOBILE) + == "ch_ppocr_mobile_v2.0_cls_mobile" + ) + + +def test_aliases_are_scoped_to_routes(): + assert ( + route_to_model_key(TaskType.REC, OCRVersion.PPOCRV4, "ja", ModelType.MOBILE) + == "japan_PP-OCRv4_rec_mobile" + ) + + +def test_list_supported_langs_includes_cls(): + assert list_supported_langs(TaskType.CLS, OCRVersion.PPOCRV5, ModelType.SERVER) == [ + "ch", + "multi", + ] diff --git a/python/tests/test_tensorrt.py b/python/tests/test_tensorrt.py index 70b2e2e..7fb2466 100644 --- a/python/tests/test_tensorrt.py +++ b/python/tests/test_tensorrt.py @@ -166,10 +166,37 @@ def test_tensorrt_default_cache_dir_uses_model_root_dir(monkeypatch, tmp_path): assert engine_paths == [ cfg.model_root_dir / "models" - / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" + / "en_PP-OCRv3_det_mobile_sm87_fp32_tf32unset.engine" ] +def test_tensorrt_model_aliases_share_cache_name(monkeypatch, tmp_path): + tensorrt_main = import_tensorrt_main(monkeypatch) + session = object.__new__(tensorrt_main.TRTInferSession) + base_cfg = { + "model_path": None, + "ocr_version": OCRVersion.PPOCRV6, + "task_type": TaskType.DET, + "model_type": ModelType.SMALL, + } + + canonical_cfg = AttrDict(base_cfg, lang_type="ch") + alias_cfg = AttrDict(base_cfg, lang_type="zh") + + assert session._get_model_name(canonical_cfg) == "multi_PP-OCRv6_det_small" + assert session._get_model_name(alias_cfg) == session._get_model_name(canonical_cfg) + + +def test_tensorrt_close_handles_partially_initialized_session(monkeypatch): + tensorrt_main = import_tensorrt_main(monkeypatch) + session = object.__new__(tensorrt_main.TRTInferSession) + + session.close() + session.close() + + assert session._closed is True + + def test_tensorrt_explicit_model_path_skips_download(monkeypatch, tmp_path): tensorrt_main = import_tensorrt_main(monkeypatch) @@ -209,7 +236,7 @@ def test_tensorrt_loads_cached_engine_without_rebuild(monkeypatch, tmp_path): cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"}) engine_path = ( - tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" + tmp_path / "trt_cache" / "en_PP-OCRv3_det_mobile_sm87_fp32_tf32unset.engine" ) engine_path.parent.mkdir() engine_path.write_bytes(b"cached engine") @@ -239,7 +266,7 @@ def test_tensorrt_force_rebuild_ignores_cached_engine(monkeypatch, tmp_path): engine_cfg={"cache_dir": tmp_path / "trt_cache", "force_rebuild": True}, ) engine_path = ( - tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" + tmp_path / "trt_cache" / "en_PP-OCRv3_det_mobile_sm87_fp16_tf32unset.engine" ) engine_path.parent.mkdir() engine_path.write_bytes(b"cached engine") @@ -278,7 +305,7 @@ def test_tensorrt_falls_back_to_rebuild_when_cached_engine_load_fails( cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"}) engine_path = ( - tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" + tmp_path / "trt_cache" / "en_PP-OCRv3_det_mobile_sm87_fp16_tf32unset.engine" ) engine_path.parent.mkdir() engine_path.write_bytes(b"bad cached engine") diff --git a/python/tools/prepare_wheel_assets.py b/python/tools/prepare_wheel_assets.py index 905be56..81cee31 100644 --- a/python/tools/prepare_wheel_assets.py +++ b/python/tools/prepare_wheel_assets.py @@ -12,7 +12,7 @@ from typing import Iterable, List, Optional PYTHON_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PYTHON_ROOT)) -from rapidocr.utils.model_resolver import normalize_lang, resolve_model_key +from rapidocr.utils.model_resolver import normalize_lang, route_to_model_key from rapidocr.utils.typings import ModelType, OCRVersion, TaskType DEFAULT_CONFIG_YAML = PYTHON_ROOT / "rapidocr" / "config.yaml" @@ -152,7 +152,7 @@ def select_model_info(spec: ModelSpec, registry: dict) -> dict: f"{spec.engine}.{spec.ocr_version}.{spec.task}" ) from e - model_key = _resolve_model_key(spec) + model_key = _route_to_model_key(spec) if model_key is not None: model_info = task_models.get(model_key) if not isinstance(model_info, dict): @@ -185,7 +185,7 @@ def select_model_info(spec: ModelSpec, registry: dict) -> dict: ) -def _resolve_model_key(spec: ModelSpec) -> Optional[str]: +def _route_to_model_key(spec: ModelSpec) -> Optional[str]: try: task_type = TaskType(spec.task) ocr_version = OCRVersion(spec.ocr_version) @@ -193,7 +193,7 @@ def _resolve_model_key(spec: ModelSpec) -> Optional[str]: except ValueError: return None - return resolve_model_key(task_type, ocr_version, spec.lang, model_type) + return route_to_model_key(task_type, ocr_version, spec.lang, model_type) def model_info_to_assets(spec: ModelSpec, model_info: dict) -> List[WheelAsset]: @@ -354,7 +354,7 @@ def _get_content_length(value: Optional[str]) -> Optional[int]: def _print_progress(downloaded: int, total_size: Optional[int]) -> None: if not total_size: - return + return percent = downloaded / total_size * 100 print( f"\r {downloaded / 1024 / 1024:.1f}MB / " diff --git a/python/tools/validate_model_routes.py b/python/tools/validate_model_routes.py new file mode 100644 index 0000000..33631d3 --- /dev/null +++ b/python/tools/validate_model_routes.py @@ -0,0 +1,108 @@ +# -*- encoding: utf-8 -*- +from collections.abc import Mapping, Sequence +from typing import Any + +from omegaconf import OmegaConf +from rapidocr.utils.typings import ModelType, OCRVersion, TaskType + + +def validate_model_routes(model_routes: Any) -> None: + """Validate the engine-neutral route schema explicitly.""" + if not _is_mapping(model_routes): + raise ValueError("Invalid model_routes: expected a mapping") + + for path, routes in _iter_model_route_lists(model_routes): + _validate_model_route_list(path, routes) + + +def _iter_model_route_lists(model_routes: Any): + for version, version_cfg in model_routes.items(): + if version == "_defs": + continue + _validate_enum("OCR version", version, OCRVersion) + if not _is_mapping(version_cfg): + raise ValueError(f"Invalid model route version: {version}") + if "aliases" in version_cfg: + _validate_aliases(f"{version}.aliases", version_cfg["aliases"]) + + for task, task_cfg in version_cfg.items(): + if task == "aliases": + continue + _validate_enum("task", task, TaskType) + if not _is_mapping(task_cfg): + raise ValueError(f"Invalid model route task: {version}.{task}") + if "aliases" in task_cfg: + _validate_aliases(f"{version}.{task}.aliases", task_cfg["aliases"]) + + for model_type, routes in task_cfg.items(): + if model_type != "aliases": + _validate_enum("model type", model_type, ModelType) + yield f"{version}.{task}.{model_type}", routes + + +def _validate_enum(label: str, value: str, enum_type: Any) -> None: + try: + enum_type(value) + except ValueError as exc: + valid_values = ", ".join(item.value for item in enum_type) + raise ValueError( + f"Invalid {label} {value!r}; expected one of: {valid_values}" + ) from exc + + +def _is_mapping(value: Any) -> bool: + return OmegaConf.is_dict(value) or isinstance(value, Mapping) + + +def _is_sequence(value: Any) -> bool: + return OmegaConf.is_list(value) or ( + isinstance(value, Sequence) and not isinstance(value, (str, bytes)) + ) + + +def _validate_aliases(path: str, aliases: Any) -> None: + if not _is_mapping(aliases): + raise ValueError(f"Invalid aliases: {path}") + + for alias, target in aliases.items(): + if not isinstance(alias, str) or not alias.strip(): + raise ValueError(f"Invalid alias: {path}") + if not isinstance(target, str) or not target.strip(): + raise ValueError(f"Invalid alias target: {path}.{alias}") + + +def _validate_model_route_list(path: str, routes: Any) -> None: + if not _is_sequence(routes): + raise ValueError(f"Invalid model route list: {path}") + + seen = set() + seen_languages = set() + for route in routes: + if not _is_mapping(route): + raise ValueError(f"Invalid model route entry: {path}") + + required = {"lang", "model_key", "supported_langs"} + if not required.issubset(route): + raise ValueError(f"Invalid model route entry: {path}") + + if not isinstance(route["lang"], str) or not route["lang"].strip(): + raise ValueError(f"Invalid route language: {path}") + + model_key = route["model_key"] + if not isinstance(model_key, str) or not model_key.strip(): + raise ValueError(f"Invalid route model key: {path}") + if model_key in seen: + raise ValueError(f"Duplicate model route key: {model_key}") + + seen.add(model_key) + + supported_langs = route["supported_langs"] + if not _is_sequence(supported_langs) or not supported_langs: + raise ValueError(f"Empty supported languages: {model_key}") + + for lang in supported_langs: + if not isinstance(lang, str) or not lang.strip(): + raise ValueError(f"Invalid supported language: {model_key}") + if lang in seen_languages: + raise ValueError(f"Overlapping supported language: {path}.{lang}") + seen_languages.add(lang)