mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-13 22:40:53 +08:00
refactor: change model cache name
This commit is contained in:
1 parent
03f3e5dd02
commit
cbc8f0d102
5 files changed
+105
-24
No files matched your search
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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')."""
|
||||
|
||||
@@ -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",
|
||||
[
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
Reference in new issue
Block a user