36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import pytest
|
|
|
|
from ocr_app.cli import build_input_parser
|
|
from ocr_app.result_adapter import model_names_for_size
|
|
from ocr_app.runtime import RuntimeConfig
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_size", "detection_model", "recognition_model"),
|
|
[
|
|
("tiny", "PP-OCRv6_tiny_det", "PP-OCRv6_tiny_rec"),
|
|
("small", "PP-OCRv6_small_det", "PP-OCRv6_small_rec"),
|
|
("medium", "PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),
|
|
],
|
|
)
|
|
def test_model_size_mapping(model_size, detection_model, recognition_model):
|
|
assert model_names_for_size(model_size) == (detection_model, recognition_model)
|
|
config = RuntimeConfig(device="cpu", model_size=model_size)
|
|
assert config.text_detection_model_name == detection_model
|
|
assert config.text_recognition_model_name == recognition_model
|
|
|
|
|
|
def test_model_size_defaults_to_medium():
|
|
args = build_input_parser().parse_args(["sample.png"])
|
|
assert args.model_size == "medium"
|
|
|
|
|
|
def test_model_alias_selects_tiny():
|
|
args = build_input_parser().parse_args(["sample.png", "--model", "tiny"])
|
|
assert args.model_size == "tiny"
|
|
|
|
|
|
def test_invalid_model_size_is_rejected():
|
|
with pytest.raises(ValueError):
|
|
model_names_for_size("large")
|