diff --git a/python/rapidocr/cli.py b/python/rapidocr/cli.py index fc5b6c3..df11ef6 100644 --- a/python/rapidocr/cli.py +++ b/python/rapidocr/cli.py @@ -3,9 +3,15 @@ # @Contact: liekkaskono@163.com import shutil from pathlib import Path +from typing import List root_dir = Path(__file__).resolve().parent DEFAULT_CFG_PATH = root_dir / "config.yaml" +REQUIRED_PACKAGE_FILES = [ + DEFAULT_CFG_PATH, + root_dir / "default_models.yaml", + root_dir / "inference_engine" / "pytorch" / "networks" / "arch_config.yaml", +] def generate_cfg(args): @@ -16,7 +22,22 @@ def generate_cfg(args): print(f"The config file has saved in {args.save_cfg_file}") -def check_install(ocr_engine): +def check_required_files() -> None: + missing_files: List[str] = [ + str(file_path) for file_path in REQUIRED_PACKAGE_FILES if not file_path.is_file() + ] + if missing_files: + missing_info = "\n".join(f" - {file_path}" for file_path in missing_files) + raise FileNotFoundError( + "RapidOCR installation is missing required package files:\n" + f"{missing_info}" + ) + + +def check_install(ocr_engine, check_files: bool = True): + if check_files: + check_required_files() + img_url = "https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.1.0/resources/test_files/ch_en_num.jpg" result = ocr_engine(img_url) diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index 979eb5a..0da2bfe 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -14,7 +14,7 @@ from .cal_rec_boxes import CalRecBoxes from .ch_ppocr_cls import TextClassifier, TextClsOutput from .ch_ppocr_det import TextDetector, TextDetOutput from .ch_ppocr_rec import TextRecInput, TextRecognizer, TextRecOutput -from .cli import check_install, generate_cfg +from .cli import check_install, check_required_files, generate_cfg from .utils.download_models import download_models from .utils.load_image import LoadImage from .utils.log import logger @@ -430,6 +430,9 @@ def main(arg_list: Optional[List[str]] = None): download_models(args.config) return + if args.command == "check": + check_required_files() + params = { "Global.text_score": args.text_score, "Global.return_word_box": args.return_word_box, @@ -438,7 +441,7 @@ def main(arg_list: Optional[List[str]] = None): ocr_engine = RapidOCR(params=params) if args.command == "check": - check_install(ocr_engine) + check_install(ocr_engine, check_files=False) return if args.img_path is None: diff --git a/python/setup.py b/python/setup.py index fb0339d..cd015de 100644 --- a/python/setup.py +++ b/python/setup.py @@ -8,6 +8,7 @@ from typing import List, Union import setuptools from get_pypi_latest_version import GetPyPiLatestVersion + from rapidocr.utils.download_models import download_models @@ -26,10 +27,19 @@ def get_readme(): return readme +def collect_package_files(package_root: Path, pattern: str) -> List[str]: + return sorted( + path.relative_to(package_root).as_posix() + for path in package_root.rglob(pattern) + if path.is_file() + ) + + download_models() MODULE_NAME = "rapidocr" PACKAGE_DIR = Path(MODULE_NAME) +PACKAGE_ROOT = PACKAGE_DIR WRAP_DIR = Path(f"wrap_temp/{MODULE_NAME}") NEED_WRAP = "bdist_wheel" in sys.argv @@ -45,6 +55,16 @@ if NEED_WRAP: f.write("from .rapidocr.main import RapidOCR, VisRes\n") PACKAGE_DIR = WRAP_DIR + PACKAGE_ROOT = PACKAGE_DIR / MODULE_NAME + +PACKAGE_DATA = collect_package_files(PACKAGE_ROOT, "*.yaml") +PACKAGE_DATA.extend( + [ + "models/ch_PP-OCRv4_det_mobile.onnx", + "models/ch_PP-OCRv4_rec_mobile.onnx", + "models/ch_ppocr_mobile_v2.0_cls_mobile.onnx", + ] +) obtainer = GetPyPiLatestVersion() try: @@ -81,14 +101,7 @@ setuptools.setup( install_requires=read_txt("requirements.txt"), package_dir={"": str(PACKAGE_DIR)}, packages=setuptools.find_namespace_packages(where=str(PACKAGE_DIR)), - package_data={ - MODULE_NAME: [ - "*.yaml", - "models/ch_PP-OCRv4_det_mobile.onnx", - "models/ch_PP-OCRv4_rec_mobile.onnx", - "models/ch_ppocr_mobile_v2.0_cls_mobile.onnx", - ] - }, + package_data={MODULE_NAME: PACKAGE_DATA}, keywords=[ "ocr,text_detection,text_recognition,db,onnxruntime,paddleocr,openvino,rapidocr" ], diff --git a/python/tests/test_cli.py b/python/tests/test_cli.py index fa4e492..f5ca3ca 100644 --- a/python/tests/test_cli.py +++ b/python/tests/test_cli.py @@ -12,6 +12,8 @@ from pytest import mark root_dir = Path(__file__).resolve().parent.parent sys.path.append(str(root_dir)) +import rapidocr.cli as cli_module +import rapidocr.main as main_module from rapidocr import LangRec, RapidOCR from rapidocr.utils.output import RapidOCROutput from rapidocr.main import main @@ -57,13 +59,44 @@ def test_cli_config(capsys, cmd): @mark.parametrize("cmd", ["check"]) -def test_cli_check(capsys, cmd): +def test_cli_check(monkeypatch, capsys, cmd): + class DummyRapidOCR: + def __init__(self, params=None): + pass + + def __call__(self, img_url): + return RapidOCROutput(txts=("正品促销",)) + + monkeypatch.setattr(main_module, "RapidOCR", DummyRapidOCR) + main(shlex.split(cmd)) output = capsys.readouterr().out.strip() assert "Success! rapidocr is installed correctly!" in output +def test_check_required_files_missing(monkeypatch, tmp_path): + missing_file = tmp_path / "missing.yaml" + monkeypatch.setattr(cli_module, "REQUIRED_PACKAGE_FILES", [missing_file]) + + with pytest.raises(FileNotFoundError, match="missing required package files"): + cli_module.check_required_files() + + +def test_cli_check_validates_required_files_before_engine(monkeypatch): + def raise_missing_file(): + raise FileNotFoundError("missing package file") + + def init_engine(*args, **kwargs): + raise AssertionError("RapidOCR should not be initialized") + + monkeypatch.setattr(main_module, "check_required_files", raise_missing_file) + monkeypatch.setattr(main_module, "RapidOCR", init_engine) + + with pytest.raises(FileNotFoundError, match="missing package file"): + main(["check"]) + + @mark.parametrize( "cmd,img_name", [