diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index 4e1939f..83b1fc2 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/python/rapidocr/inference_engine/base.py @@ -140,15 +140,6 @@ class InferSession(abc.ABC): if model_key in model_dict: return model_dict[model_key] - # Keep legacy Torch classification registry naming compatible with - # the engine-neutral route key used by the shared model registry. - legacy_keys = { - "ch_ppocr_mobile_v2.0_cls_mobile": "ch_ptocr_mobile_v2.0_cls_mobile", - } - legacy_key = legacy_keys.get(model_key) - if legacy_key in model_dict: - return model_dict[legacy_key] - raise ValueError( f"Unsupported configuration: {engine_type}.{ocr_version}.{task_type}.{lang_type}.{model_type}" ) diff --git a/python/rapidocr/inference_engine/tensorrt/main.py b/python/rapidocr/inference_engine/tensorrt/main.py index 811ecd1..d294b5e 100644 --- a/python/rapidocr/inference_engine/tensorrt/main.py +++ b/python/rapidocr/inference_engine/tensorrt/main.py @@ -234,15 +234,18 @@ class TRTInferSession(InferSession): if cfg.get("model_path"): return Path(cfg["model_path"]).stem - 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 - # Cache names identify the requested configuration. Some PP-OCRv4 - # routes intentionally reuse PP-OCRv3 assets (for example English - # detection), but their TensorRT cache must remain scoped to the - # configured OCR version to avoid collisions. - return f"{lang_type}_{ocr_version}_{task_type}_{model_type}" + model_key = route_to_model_key( + cfg.task_type, + cfg.ocr_version, + cfg.lang_type, + cfg.model_type, + ) + if model_key is None: + task_type = cfg.task_type.value + lang_type = normalize_lang(cfg.lang_type) + model_key = f"{lang_type}_{task_type}_{cfg.model_type.value}" + + return model_key def _get_gpu_arch(self) -> str: """Get GPU architecture string for cache key (e.g., 'sm87').""" diff --git a/python/tests/test_model_resolver.py b/python/tests/test_model_resolver.py index 88f2a76..10b92ae 100644 --- a/python/tests/test_model_resolver.py +++ b/python/tests/test_model_resolver.py @@ -2,8 +2,8 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import pytest - from rapidocr import ModelType, OCRVersion +from rapidocr.inference_engine.base import InferSession from rapidocr.utils.model_resolver import ( MODEL_ROUTES, list_supported_langs, @@ -17,6 +17,37 @@ def test_model_routes_schema_is_valid(): validate_model_routes(MODEL_ROUTES) +def test_every_route_resolves_and_exists_in_a_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"] + assert any( + model_key in (registry.get(version, {}).get(task, {}) or {}) + for engine, registry in registries.items() + if engine != "model_routes" + ), model_key + for lang in route["supported_langs"]: + assert ( + route_to_model_key( + task_enum, version_enum, lang, model_enum + ) + == model_key + ) + + @pytest.mark.parametrize( "task_type,expected_key", [ diff --git a/python/tests/test_tensorrt.py b/python/tests/test_tensorrt.py index 70b2e2e..f34cd26 100644 --- a/python/tests/test_tensorrt.py +++ b/python/tests/test_tensorrt.py @@ -166,10 +166,27 @@ 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_fp16_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_explicit_model_path_skips_download(monkeypatch, tmp_path): tensorrt_main = import_tensorrt_main(monkeypatch) @@ -209,7 +226,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_fp16_tf32unset.engine" ) engine_path.parent.mkdir() engine_path.write_bytes(b"cached engine") @@ -239,7 +256,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 +295,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/validate_model_routes.py b/python/tools/validate_model_routes.py index 2af00c2..0c248b3 100644 --- a/python/tools/validate_model_routes.py +++ b/python/tools/validate_model_routes.py @@ -6,6 +6,9 @@ from omegaconf import OmegaConf def validate_model_routes(model_routes: Any) -> None: """Validate the engine-neutral route schema explicitly.""" + if not OmegaConf.is_dict(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) @@ -14,31 +17,67 @@ def _iter_model_route_lists(model_routes: Any): for version, version_cfg in model_routes.items(): if version == "_defs": continue + if not OmegaConf.is_dict(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 + if not OmegaConf.is_dict(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": yield f"{version}.{task}.{model_type}", routes +def _validate_aliases(path: str, aliases: Any) -> None: + if not OmegaConf.is_dict(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 OmegaConf.is_list(routes): raise ValueError(f"Invalid model route list: {path}") seen = set() + seen_languages = set() for route in routes: + if not OmegaConf.is_dict(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) - if not route["supported_langs"]: + supported_langs = route["supported_langs"] + if not OmegaConf.is_list(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)