mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 05:20:49 +08:00
297 lines
10 KiB
Python
297 lines
10 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from rapidocr import RapidOCR
|
|
from rapidocr.inference_engine.base import FileInfo, InferSession
|
|
from rapidocr.utils.parse_parameters import ParseParams
|
|
from rapidocr.utils.typings import (
|
|
EngineType,
|
|
LangCls,
|
|
LangDet,
|
|
LangRec,
|
|
ModelType,
|
|
OCRVersion,
|
|
TaskType,
|
|
)
|
|
|
|
root_dir = Path(__file__).resolve().parent.parent
|
|
config_path = root_dir / "rapidocr" / "config.yaml"
|
|
tests_dir = root_dir / "tests" / "test_files"
|
|
|
|
|
|
@pytest.fixture()
|
|
def default_cfg():
|
|
return ParseParams.load(config_path)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def default_engine():
|
|
return RapidOCR()
|
|
|
|
|
|
def get_default_model_info(model_cfg):
|
|
return InferSession.get_model_url(
|
|
FileInfo(
|
|
model_cfg.engine_type,
|
|
model_cfg.ocr_version,
|
|
model_cfg.task_type,
|
|
model_cfg.lang_type,
|
|
model_cfg.model_type,
|
|
)
|
|
)
|
|
|
|
|
|
def test_default_global_config(default_cfg):
|
|
assert default_cfg.Global.text_score == 0.5
|
|
assert default_cfg.Global.use_det is True
|
|
assert default_cfg.Global.use_cls is True
|
|
assert default_cfg.Global.use_rec is True
|
|
assert default_cfg.Global.min_height == 30
|
|
assert default_cfg.Global.width_height_ratio == 8
|
|
assert default_cfg.Global.max_side_len == 2000
|
|
assert default_cfg.Global.min_side_len == 30
|
|
assert default_cfg.Global.return_word_box is False
|
|
assert default_cfg.Global.return_single_char_box is False
|
|
assert default_cfg.Global.font_path is None
|
|
assert default_cfg.Global.log_level == "info"
|
|
assert default_cfg.Global.model_root_dir is None
|
|
|
|
|
|
def test_default_onnxruntime_engine_config(default_cfg):
|
|
cfg = default_cfg.EngineConfig.onnxruntime
|
|
|
|
assert cfg.intra_op_num_threads == -1
|
|
assert cfg.inter_op_num_threads == -1
|
|
assert cfg.enable_cpu_mem_arena is False
|
|
assert cfg.cpu_ep_cfg.arena_extend_strategy == "kSameAsRequested"
|
|
assert cfg.use_cuda is False
|
|
assert cfg.cuda_ep_cfg.device_id == 0
|
|
assert cfg.cuda_ep_cfg.arena_extend_strategy == "kNextPowerOfTwo"
|
|
assert cfg.cuda_ep_cfg.cudnn_conv_algo_search == "EXHAUSTIVE"
|
|
assert cfg.cuda_ep_cfg.do_copy_in_default_stream is True
|
|
assert cfg.use_dml is False
|
|
assert cfg.dml_ep_cfg is None
|
|
assert cfg.use_cann is False
|
|
assert cfg.cann_ep_cfg.device_id == 0
|
|
assert cfg.cann_ep_cfg.npu_mem_limit == 21474836480
|
|
assert cfg.cann_ep_cfg.op_select_impl_mode == "high_performance"
|
|
assert cfg.cann_ep_cfg.optypelist_for_implmode == "Gelu"
|
|
assert cfg.cann_ep_cfg.enable_cann_graph is True
|
|
assert cfg.use_coreml is False
|
|
assert cfg.coreml_ep_cfg.ModelFormat == "MLProgram"
|
|
assert cfg.coreml_ep_cfg.MLComputeUnits == "ALL"
|
|
assert cfg.coreml_ep_cfg.ModelCacheDirectory == "/tmp/RapidOCR"
|
|
|
|
|
|
def test_default_openvino_engine_config(default_cfg):
|
|
cfg = default_cfg.EngineConfig.openvino
|
|
|
|
assert cfg.inference_num_threads == -1
|
|
assert cfg.performance_hint is None
|
|
assert cfg.performance_num_requests == -1
|
|
assert cfg.enable_cpu_pinning is None
|
|
assert cfg.num_streams == -1
|
|
assert cfg.enable_hyper_threading is None
|
|
assert cfg.scheduling_core_type is None
|
|
|
|
|
|
def test_default_other_engine_configs(default_cfg):
|
|
assert default_cfg.EngineConfig.paddle.cpu_math_library_num_threads == -1
|
|
assert default_cfg.EngineConfig.paddle.use_npu is False
|
|
assert default_cfg.EngineConfig.paddle.npu_ep_cfg.device_id == 0
|
|
assert default_cfg.EngineConfig.paddle.use_cuda is False
|
|
assert default_cfg.EngineConfig.paddle.cuda_ep_cfg.device_id == 0
|
|
assert default_cfg.EngineConfig.paddle.cuda_ep_cfg.gpu_mem == 500
|
|
|
|
assert default_cfg.EngineConfig.torch.use_cuda is False
|
|
assert default_cfg.EngineConfig.torch.cuda_ep_cfg.device_id == 0
|
|
assert default_cfg.EngineConfig.torch.use_npu is False
|
|
assert default_cfg.EngineConfig.torch.npu_ep_cfg.device_id == 0
|
|
assert default_cfg.EngineConfig.torch.use_mps is False
|
|
|
|
assert default_cfg.EngineConfig.tensorrt.device_id == 0
|
|
assert default_cfg.EngineConfig.tensorrt.use_fp16 is True
|
|
assert default_cfg.EngineConfig.tensorrt.use_int8 is False
|
|
assert default_cfg.EngineConfig.tensorrt.workspace_size == 1073741824
|
|
assert default_cfg.EngineConfig.tensorrt.cache_dir is None
|
|
assert default_cfg.EngineConfig.tensorrt.force_rebuild is False
|
|
assert default_cfg.EngineConfig.tensorrt.det_profile.opt_shape == [
|
|
1,
|
|
3,
|
|
736,
|
|
736,
|
|
]
|
|
assert default_cfg.EngineConfig.tensorrt.rec_profile.opt_shape == [
|
|
6,
|
|
3,
|
|
48,
|
|
320,
|
|
]
|
|
assert default_cfg.EngineConfig.tensorrt.cls_profile.opt_shape == [
|
|
6,
|
|
3,
|
|
48,
|
|
192,
|
|
]
|
|
|
|
assert default_cfg.EngineConfig.mnn == {}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"section,task_type,expected_model_name,expected_sha256",
|
|
[
|
|
(
|
|
"Det",
|
|
TaskType.DET,
|
|
"PP-OCRv6_det_small.onnx",
|
|
"090f04abcd9d9a7498bc4ebf677e4cb9bdce1fe4197ddb7e529f1ef44e1ff94f",
|
|
),
|
|
(
|
|
"Rec",
|
|
TaskType.REC,
|
|
"PP-OCRv6_rec_small.onnx",
|
|
"6f327246b50388f3c176ae304bd95767ea6dc0c9ae92153ef8cbe210b3c14884",
|
|
),
|
|
],
|
|
)
|
|
def test_default_ppocrv6_model_config(
|
|
default_cfg, section, task_type, expected_model_name, expected_sha256
|
|
):
|
|
model_cfg = default_cfg[section]
|
|
|
|
assert model_cfg.engine_type == EngineType.ONNXRUNTIME
|
|
assert model_cfg.ocr_version == OCRVersion.PPOCRV6
|
|
assert model_cfg.model_type == ModelType.SMALL
|
|
assert model_cfg.task_type == task_type
|
|
|
|
model_info = get_default_model_info(model_cfg)
|
|
|
|
assert model_info["model_dir"].endswith(
|
|
f"/PP-OCRv6/{task_type.value}/{expected_model_name}"
|
|
)
|
|
assert model_info["SHA256"] == expected_sha256
|
|
|
|
|
|
def test_default_det_config(default_cfg):
|
|
assert default_cfg.Det.engine_type == EngineType.ONNXRUNTIME
|
|
assert default_cfg.Det.lang_type == LangDet.CH
|
|
assert default_cfg.Det.model_type == ModelType.SMALL
|
|
assert default_cfg.Det.ocr_version == OCRVersion.PPOCRV6
|
|
assert default_cfg.Det.task_type == TaskType.DET
|
|
assert default_cfg.Det.model_path is None
|
|
assert default_cfg.Det.model_dir is None
|
|
assert default_cfg.Det.limit_side_len == 736
|
|
assert default_cfg.Det.limit_type == "min"
|
|
assert default_cfg.Det.std == [0.5, 0.5, 0.5]
|
|
assert default_cfg.Det.mean == [0.5, 0.5, 0.5]
|
|
assert default_cfg.Det.thresh == 0.3
|
|
assert default_cfg.Det.box_thresh == 0.5
|
|
assert default_cfg.Det.max_candidates == 1000
|
|
assert default_cfg.Det.unclip_ratio == 1.6
|
|
assert default_cfg.Det.use_dilation is True
|
|
assert default_cfg.Det.score_mode == "fast"
|
|
|
|
|
|
def test_default_cls_config(default_cfg):
|
|
assert default_cfg.Cls.engine_type == EngineType.ONNXRUNTIME
|
|
assert default_cfg.Cls.lang_type == LangCls.CH
|
|
assert default_cfg.Cls.model_type == ModelType.MOBILE
|
|
assert default_cfg.Cls.ocr_version == OCRVersion.PPOCRV4
|
|
assert default_cfg.Cls.task_type == TaskType.CLS
|
|
assert default_cfg.Cls.model_path is None
|
|
assert default_cfg.Cls.model_dir is None
|
|
assert default_cfg.Cls.cls_image_shape == [3, 48, 192]
|
|
assert default_cfg.Cls.cls_batch_num == 6
|
|
assert default_cfg.Cls.cls_thresh == 0.9
|
|
assert default_cfg.Cls.label_list == ["0", "180"]
|
|
|
|
|
|
def test_default_rec_config(default_cfg):
|
|
assert default_cfg.Rec.engine_type == EngineType.ONNXRUNTIME
|
|
assert default_cfg.Rec.lang_type == LangRec.CH
|
|
assert default_cfg.Rec.model_type == ModelType.SMALL
|
|
assert default_cfg.Rec.ocr_version == OCRVersion.PPOCRV6
|
|
assert default_cfg.Rec.task_type == TaskType.REC
|
|
assert default_cfg.Rec.model_path is None
|
|
assert default_cfg.Rec.model_dir is None
|
|
assert default_cfg.Rec.rec_keys_path is None
|
|
assert default_cfg.Rec.rec_img_shape == [3, 48, 320]
|
|
assert default_cfg.Rec.rec_batch_num == 6
|
|
|
|
|
|
def test_default_cls_model_config(default_cfg):
|
|
model_info = get_default_model_info(default_cfg.Cls)
|
|
|
|
assert model_info["model_dir"].endswith(
|
|
"/PP-OCRv4/cls/ch_ppocr_mobile_v2.0_cls_mobile.onnx"
|
|
)
|
|
assert (
|
|
model_info["SHA256"]
|
|
== "e47acedf663230f8863ff1ab0e64dd2d82b838fceb5957146dab185a89d6215c"
|
|
)
|
|
|
|
|
|
def test_default_engine_runtime_config(default_engine):
|
|
assert default_engine.use_det is True
|
|
assert default_engine.use_cls is True
|
|
assert default_engine.use_rec is True
|
|
assert default_engine.cfg.Det.ocr_version == OCRVersion.PPOCRV6
|
|
assert default_engine.cfg.Det.model_type == ModelType.SMALL
|
|
assert default_engine.cfg.Rec.ocr_version == OCRVersion.PPOCRV6
|
|
assert default_engine.cfg.Rec.model_type == ModelType.SMALL
|
|
assert default_engine.cfg.Cls.ocr_version == OCRVersion.PPOCRV4
|
|
assert default_engine.cfg.Cls.model_type == ModelType.MOBILE
|
|
|
|
|
|
def test_update_det_params_before_model_load():
|
|
engine = RapidOCR()
|
|
|
|
assert engine.text_det is None
|
|
|
|
engine.update_params(box_thresh=0.7, unclip_ratio=2.0)
|
|
|
|
assert engine.text_det is None
|
|
assert engine.cfg.Det.box_thresh == 0.7
|
|
assert engine.cfg.Det.unclip_ratio == 2.0
|
|
|
|
|
|
def test_update_det_params_after_model_load():
|
|
engine = RapidOCR()
|
|
engine.text_det = SimpleNamespace(
|
|
postprocess_op=SimpleNamespace(box_thresh=0.5, unclip_ratio=1.6)
|
|
)
|
|
|
|
engine.update_params(box_thresh=0.7, unclip_ratio=2.0)
|
|
|
|
assert engine.cfg.Det.box_thresh == 0.7
|
|
assert engine.cfg.Det.unclip_ratio == 2.0
|
|
assert engine.text_det.postprocess_op.box_thresh == 0.7
|
|
assert engine.text_det.postprocess_op.unclip_ratio == 2.0
|
|
|
|
|
|
def test_default_engine_real_full_pipeline(default_engine):
|
|
result = default_engine(tests_dir / "ch_en_num.jpg")
|
|
|
|
assert result.txts is not None
|
|
assert len(result) == 18
|
|
assert result.txts[:3] == ("正品促销", "大桶装更划算", "强力去污符合国标")
|
|
|
|
|
|
def test_default_engine_real_rec_only(default_engine):
|
|
result = default_engine(
|
|
tests_dir / "text_rec.jpg",
|
|
use_det=False,
|
|
use_cls=False,
|
|
use_rec=True,
|
|
)
|
|
|
|
assert result.txts is not None
|
|
assert len(result) == 1
|
|
assert result.txts[0] == "韩国小馆"
|
|
assert result.scores[0] > 0.99
|