mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-10 21:10:55 +08:00
71 lines
1.7 KiB
Python
71 lines
1.7 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
root_dir = Path(__file__).resolve().parent.parent
|
|
sys.path.append(str(root_dir))
|
|
|
|
from tools.prepare_wheel_assets import (
|
|
ModelSpec,
|
|
load_yaml,
|
|
resolve_assets,
|
|
resolve_model_specs,
|
|
select_model_info,
|
|
)
|
|
|
|
config_path = root_dir / "rapidocr" / "config.yaml"
|
|
models_path = root_dir / "rapidocr" / "default_models.yaml"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def default_config():
|
|
return load_yaml(config_path)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def model_registry():
|
|
return load_yaml(models_path)
|
|
|
|
|
|
def test_default_wheel_assets_use_routed_ppocrv6_models(
|
|
default_config, model_registry
|
|
):
|
|
specs = resolve_model_specs(default_config)
|
|
assets = resolve_assets(specs, model_registry)
|
|
asset_paths = {asset.relative_path.as_posix() for asset in assets}
|
|
|
|
assert asset_paths == {
|
|
"PP-OCRv6_det_small.onnx",
|
|
"ch_ppocr_mobile_v2.0_cls_mobile.onnx",
|
|
"PP-OCRv6_rec_small.onnx",
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"spec,expected_model_name",
|
|
[
|
|
(
|
|
ModelSpec("onnxruntime", "PP-OCRv6", "det", "ch", "small"),
|
|
"PP-OCRv6_det_small.onnx",
|
|
),
|
|
(
|
|
ModelSpec("onnxruntime", "PP-OCRv6", "rec", "ch", "small"),
|
|
"PP-OCRv6_rec_small.onnx",
|
|
),
|
|
(
|
|
ModelSpec("onnxruntime", "PP-OCRv4", "cls", "ch", "mobile"),
|
|
"ch_ppocr_mobile_v2.0_cls_mobile.onnx",
|
|
),
|
|
],
|
|
)
|
|
def test_select_model_info_supports_default_model_specs(
|
|
model_registry, spec, expected_model_name
|
|
):
|
|
model_info = select_model_info(spec, model_registry)
|
|
|
|
assert model_info["model_dir"].endswith(expected_model_name)
|