first commit

tests/pipelines/test_doc_preprocessor.py

@@ -0,0 +1,57 @@
+import pytest
+
+from paddleocr import DocPreprocessor
+from ..testing_utils import (
+    TEST_DATA_DIR,
+    check_simple_inference_result,
+    check_wrapper_simple_inference_param_forwarding,
+)
+
+
+@pytest.fixture(scope="module")
+def ocr_engine() -> DocPreprocessor:
+    return DocPreprocessor()
+
+
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "book_rot180.jpg",
+    ],
+)
+def test_predict(ocr_engine: DocPreprocessor, image_path: str) -> None:
+    """
+    Test PaddleOCR's doc preprocessor functionality.
+
+    Args:
+        ocr_engine: An instance of `DocPreprocessor`.
+        image_path: Path to the image to be processed.
+    """
+    result = ocr_engine.predict(str(image_path))
+
+    check_simple_inference_result(result)
+    res = result[0]
+    assert res["angle"] in {0, 90, 180, 270, -1}
+    assert res["rot_img"] is not None
+    assert res["output_img"] is not None
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False},
+        {"use_doc_unwarping": False},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    ocr_engine: DocPreprocessor,
+    params: dict,
+) -> None:
+    check_wrapper_simple_inference_param_forwarding(
+        monkeypatch,
+        ocr_engine,
+        "paddlex_pipeline",
+        "dummy_path",
+        params,
+    )

tests/pipelines/test_doc_understanding.py

@@ -0,0 +1,43 @@
+import pytest
+
+from paddleocr import DocUnderstanding
+from ..testing_utils import (
+    TEST_DATA_DIR,
+    check_simple_inference_result,
+    check_wrapper_simple_inference_param_forwarding,
+)
+
+
+@pytest.fixture(scope="module")
+def ocr_engine() -> DocUnderstanding:
+    return DocUnderstanding()
+
+
+@pytest.mark.resource_intensive
+@pytest.mark.parametrize(
+    "input",
+    [
+        {
+            "image": str(TEST_DATA_DIR / "medal_table.png"),
+            "query": "识别这份表格的内容",
+        },
+        {
+            "image": str(TEST_DATA_DIR / "table.jpg"),
+            "query": "识别这份表格的内容",
+        },
+    ],
+)
+def test_predict(ocr_engine: DocUnderstanding, input: dict) -> None:
+    """
+    Test PaddleOCR's doc understanding functionality.
+
+    Args:
+        ocr_engine: An instance of `DocUnderstanding`.
+        input: Input dict to be processed.
+    """
+    result = ocr_engine.predict(input)
+
+    check_simple_inference_result(result)
+    res = result[0]
+    assert res["result"] is not None
+    assert isinstance(res["result"], str)

tests/pipelines/test_formula_recognition.py

@@ -0,0 +1,68 @@
+import pytest
+
+from paddleocr import FormulaRecognitionPipeline
+from ..testing_utils import (
+    TEST_DATA_DIR,
+    check_simple_inference_result,
+    check_wrapper_simple_inference_param_forwarding,
+)
+
+
+@pytest.fixture(scope="module")
+def formula_recognition_engine() -> FormulaRecognitionPipeline:
+    return FormulaRecognitionPipeline()
+
+
+# TODO: Should we separate unit tests and integration tests?
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "doc_with_formula.png",
+    ],
+)
+def test_predict(
+    formula_recognition_engine: FormulaRecognitionPipeline, image_path: str
+) -> None:
+    """
+    Test FormulaRecognitionPipeline's formula_recognition functionality.
+
+    Args:
+        formula_recognition_engine: An instance of `FormulaRecognitionPipeline`.
+        image_path: Path to the image to be processed.
+    """
+    result = formula_recognition_engine.predict(str(image_path))
+
+    check_simple_inference_result(result)
+    res = result[0]
+    assert isinstance(res["formula_res_list"], list)
+    assert len(res["formula_res_list"]) > 0
+
+
+# TODO: Also check passing `None`
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False},
+        {"use_doc_unwarping": False},
+        {"use_layout_detection": False},
+        {"layout_threshold": 0.5},
+        {"layout_nms": True},
+        {"layout_unclip_ratio": 1.5},
+        {"layout_merge_bboxes_mode": "large"},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    formula_recognition_engine: FormulaRecognitionPipeline,
+    params: dict,
+) -> None:
+    check_wrapper_simple_inference_param_forwarding(
+        monkeypatch,
+        formula_recognition_engine,
+        "paddlex_pipeline",
+        "dummy_path",
+        params,
+    )
+
+
+# TODO: Test init params

tests/pipelines/test_ocr.py

@@ -0,0 +1,125 @@
+import pytest
+
+from paddleocr import PaddleOCR
+from ..testing_utils import (
+    TEST_DATA_DIR,
+    check_simple_inference_result,
+    check_wrapper_simple_inference_param_forwarding,
+)
+
+
+@pytest.fixture(scope="module")
+def ocr_engine() -> PaddleOCR:
+    return PaddleOCR()
+
+
+# TODO: Should we separate unit tests and integration tests?
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "table.jpg",
+    ],
+)
+def test_predict(ocr_engine: PaddleOCR, image_path: str) -> None:
+    """
+    Test PaddleOCR's OCR functionality.
+
+    Args:
+        ocr_engine: An instance of `PaddleOCR`.
+        image_path: Path to the image to be processed.
+    """
+    result = ocr_engine.predict(str(image_path))
+
+    check_simple_inference_result(result)
+    res = result[0]
+    assert len(res["dt_polys"]) > 0
+    assert isinstance(res["rec_texts"], list)
+    assert len(res["rec_texts"]) > 0
+    for text in res["rec_texts"]:
+        assert isinstance(text, str)
+
+
+# TODO: Also check passing `None`
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False},
+        {"use_doc_unwarping": False},
+        {"use_textline_orientation": False},
+        {"text_det_limit_side_len": 640, "text_det_limit_type": "min"},
+        {"text_det_thresh": 0.5},
+        {"text_det_box_thresh": 0.3},
+        {"text_det_unclip_ratio": 3.0},
+        {"text_rec_score_thresh": 0.5},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    ocr_engine: PaddleOCR,
+    params: dict,
+) -> None:
+    check_wrapper_simple_inference_param_forwarding(
+        monkeypatch,
+        ocr_engine,
+        "paddlex_pipeline",
+        "dummy_path",
+        params,
+    )
+
+
+# TODO: Test init params
+
+
+def test_lang_and_ocr_version():
+    ocr_engine = PaddleOCR(lang="ch", ocr_version="PP-OCRv5")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv5_server_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "PP-OCRv5_server_rec"
+    ocr_engine = PaddleOCR(lang="chinese_cht", ocr_version="PP-OCRv5")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv5_server_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "PP-OCRv5_server_rec"
+    ocr_engine = PaddleOCR(lang="en", ocr_version="PP-OCRv5")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv5_server_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "PP-OCRv5_server_rec"
+    ocr_engine = PaddleOCR(lang="japan", ocr_version="PP-OCRv5")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv5_server_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "PP-OCRv5_server_rec"
+    ocr_engine = PaddleOCR(lang="ch", ocr_version="PP-OCRv4")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv4_mobile_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "PP-OCRv4_mobile_rec"
+    ocr_engine = PaddleOCR(lang="en", ocr_version="PP-OCRv4")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv4_mobile_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "en_PP-OCRv4_mobile_rec"
+    ocr_engine = PaddleOCR(lang="ch", ocr_version="PP-OCRv3")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv3_mobile_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "PP-OCRv3_mobile_rec"
+    ocr_engine = PaddleOCR(lang="en", ocr_version="PP-OCRv3")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv3_mobile_det"
+    assert ocr_engine._params["text_recognition_model_name"] == "en_PP-OCRv3_mobile_rec"
+    ocr_engine = PaddleOCR(lang="fr", ocr_version="PP-OCRv3")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv3_mobile_det"
+    assert (
+        ocr_engine._params["text_recognition_model_name"] == "latin_PP-OCRv3_mobile_rec"
+    )
+    ocr_engine = PaddleOCR(lang="ar", ocr_version="PP-OCRv3")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv3_mobile_det"
+    assert (
+        ocr_engine._params["text_recognition_model_name"]
+        == "arabic_PP-OCRv3_mobile_rec"
+    )
+    ocr_engine = PaddleOCR(lang="ru", ocr_version="PP-OCRv3")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv3_mobile_det"
+    assert (
+        ocr_engine._params["text_recognition_model_name"]
+        == "cyrillic_PP-OCRv3_mobile_rec"
+    )
+    ocr_engine = PaddleOCR(lang="hi", ocr_version="PP-OCRv3")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv3_mobile_det"
+    assert (
+        ocr_engine._params["text_recognition_model_name"]
+        == "devanagari_PP-OCRv3_mobile_rec"
+    )
+    ocr_engine = PaddleOCR(lang="japan", ocr_version="PP-OCRv3")
+    assert ocr_engine._params["text_detection_model_name"] == "PP-OCRv3_mobile_det"
+    assert (
+        ocr_engine._params["text_recognition_model_name"] == "japan_PP-OCRv3_mobile_rec"
+    )

tests/pipelines/test_pp_chatocrv4_doc.py

@@ -0,0 +1,80 @@
+import pytest
+
+from paddleocr import PPChatOCRv4Doc
+from ..testing_utils import TEST_DATA_DIR
+
+
+@pytest.fixture(scope="module")
+def pp_chatocrv4_doc_pipeline():
+    return PPChatOCRv4Doc()
+
+
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "doc_with_formula.png",
+    ],
+)
+def test_visual_predict(pp_chatocrv4_doc_pipeline, image_path):
+    result = pp_chatocrv4_doc_pipeline.visual_predict(str(image_path))
+
+    assert result is not None
+    assert isinstance(result, list)
+    assert len(result) == 1
+    res = result[0]
+    assert isinstance(res, dict)
+    assert res.keys() == {"visual_info", "layout_parsing_result"}
+    assert isinstance(res["visual_info"], dict)
+    assert isinstance(res["layout_parsing_result"], dict)
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False},
+        {"use_doc_unwarping": False},
+        {"use_table_recognition": False},
+        {"layout_threshold": 0.88},
+        {"layout_threshold": [0.45, 0.4]},
+        {"layout_threshold": {0: 0.45, 2: 0.48, 7: 0.4}},
+        {"layout_nms": False},
+        {"layout_unclip_ratio": 1.1},
+        {"layout_unclip_ratio": [1.2, 1.5]},
+        {"layout_unclip_ratio": {0: 1.2, 2: 1.5, 7: 1.8}},
+        {"layout_merge_bboxes_mode": "large"},
+        {"layout_merge_bboxes_mode": {0: "large", 2: "small", 7: "union"}},
+        {"text_det_limit_side_len": 640, "text_det_limit_type": "min"},
+        {"text_det_thresh": 0.5},
+        {"text_det_box_thresh": 0.3},
+        {"text_det_unclip_ratio": 3.0},
+        {"text_rec_score_thresh": 0.5},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    pp_chatocrv4_doc_pipeline,
+    params,
+):
+    def _dummy_visual_predict(input, **params):
+        yield {"visual_info": {}, "layout_parsing_result": params}
+
+    monkeypatch.setattr(
+        pp_chatocrv4_doc_pipeline.paddlex_pipeline,
+        "visual_predict",
+        _dummy_visual_predict,
+    )
+
+    result = pp_chatocrv4_doc_pipeline.visual_predict(
+        input,
+        **params,
+    )
+
+    assert isinstance(result, list)
+    assert len(result) == 1
+    res = result[0]
+    res = res["layout_parsing_result"]
+    for k, v in params.items():
+        assert res[k] == v
+
+
+# TODO: Test constructor and other methods

tests/pipelines/test_pp_doctranslation.py

@@ -0,0 +1,80 @@
+import pytest
+
+from paddleocr import PPDocTranslation
+from ..testing_utils import TEST_DATA_DIR
+
+
+@pytest.fixture(scope="module")
+def pp_doctranslation_pipeline():
+    return PPDocTranslation()
+
+
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "book.jpg",
+    ],
+)
+def test_visual_predict(pp_doctranslation_pipeline, image_path):
+    result = pp_doctranslation_pipeline.visual_predict(str(image_path))
+
+    assert result is not None
+    assert isinstance(result, list)
+    assert len(result) == 1
+    res = result[0]
+    assert isinstance(res, dict)
+    assert res.keys() == {"layout_parsing_result"}
+    assert isinstance(res["layout_parsing_result"], dict)
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False},
+        {"use_doc_unwarping": False},
+        {"use_table_recognition": False},
+        {"use_formula_recognition": False},
+        {"layout_threshold": 0.88},
+        {"layout_threshold": [0.45, 0.4]},
+        {"layout_threshold": {0: 0.45, 2: 0.48, 7: 0.4}},
+        {"layout_nms": False},
+        {"layout_unclip_ratio": 1.1},
+        {"layout_unclip_ratio": [1.2, 1.5]},
+        {"layout_unclip_ratio": {0: 1.2, 2: 1.5, 7: 1.8}},
+        {"layout_merge_bboxes_mode": "large"},
+        {"layout_merge_bboxes_mode": {0: "large", 2: "small", 7: "union"}},
+        {"text_det_limit_side_len": 640, "text_det_limit_type": "min"},
+        {"text_det_thresh": 0.5},
+        {"text_det_box_thresh": 0.3},
+        {"text_det_unclip_ratio": 3.0},
+        {"text_rec_score_thresh": 0.5},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    pp_doctranslation_pipeline,
+    params,
+):
+    def _dummy_visual_predict(input, **params):
+        yield {"layout_parsing_result": params}
+
+    monkeypatch.setattr(
+        pp_doctranslation_pipeline.paddlex_pipeline,
+        "visual_predict",
+        _dummy_visual_predict,
+    )
+
+    result = pp_doctranslation_pipeline.visual_predict(
+        input,
+        **params,
+    )
+
+    assert isinstance(result, list)
+    assert len(result) == 1
+    res = result[0]
+    res = res["layout_parsing_result"]
+    for k, v in params.items():
+        assert res[k] == v
+
+
+# TODO: Test constructor and other methods

tests/pipelines/test_pp_structurev3.py

@@ -0,0 +1,71 @@
+import pytest
+
+from paddleocr import PPStructureV3
+from ..testing_utils import (
+    TEST_DATA_DIR,
+    check_simple_inference_result,
+    check_wrapper_simple_inference_param_forwarding,
+)
+
+
+@pytest.fixture(scope="module")
+def pp_structurev3_pipeline():
+    return PPStructureV3()
+
+
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "doc_with_formula.png",
+    ],
+)
+def test_visual_predict(pp_structurev3_pipeline, image_path):
+    result = pp_structurev3_pipeline.predict(str(image_path))
+
+    check_simple_inference_result(result)
+    res = result[0]
+    overall_ocr_res = res["overall_ocr_res"]
+    assert len(overall_ocr_res["dt_polys"]) > 0
+    assert len(overall_ocr_res["rec_texts"]) > 0
+    assert len(overall_ocr_res["rec_polys"]) > 0
+    assert len(overall_ocr_res["rec_boxes"]) > 0
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False},
+        {"use_doc_unwarping": False},
+        {"use_table_recognition": False},
+        {"use_formula_recognition": False},
+        {"layout_threshold": 0.88},
+        {"layout_threshold": [0.45, 0.4]},
+        {"layout_threshold": {0: 0.45, 2: 0.48, 7: 0.4}},
+        {"layout_nms": False},
+        {"layout_unclip_ratio": 1.1},
+        {"layout_unclip_ratio": [1.2, 1.5]},
+        {"layout_unclip_ratio": {0: 1.2, 2: 1.5, 7: 1.8}},
+        {"layout_merge_bboxes_mode": "large"},
+        {"layout_merge_bboxes_mode": {0: "large", 2: "small", 7: "union"}},
+        {"text_det_limit_side_len": 640, "text_det_limit_type": "min"},
+        {"text_det_thresh": 0.5},
+        {"text_det_box_thresh": 0.3},
+        {"text_det_unclip_ratio": 3.0},
+        {"text_rec_score_thresh": 0.5},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    pp_structurev3_pipeline,
+    params,
+):
+    check_wrapper_simple_inference_param_forwarding(
+        monkeypatch,
+        pp_structurev3_pipeline,
+        "paddlex_pipeline",
+        "dummy_path",
+        params,
+    )
+
+
+# TODO: Test constructor and other methods

tests/pipelines/test_seal_rec.py

@@ -0,0 +1,70 @@
+import pytest
+
+from paddleocr import SealRecognition
+from ..testing_utils import (
+    TEST_DATA_DIR,
+    check_simple_inference_result,
+    check_wrapper_simple_inference_param_forwarding,
+)
+
+
+@pytest.fixture(scope="module")
+def ocr_engine() -> SealRecognition:
+    return SealRecognition()
+
+
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "seal.png",
+    ],
+)
+def test_predict(ocr_engine: SealRecognition, image_path: str) -> None:
+    """
+    Test PaddleOCR's seal recognition functionality.
+
+    Args:
+        ocr_engine: An instance of `SealRecognition`.
+        image_path: Path to the image to be processed.
+    """
+    result = ocr_engine.predict(str(image_path))
+
+    check_simple_inference_result(result)
+    res = result[0]["seal_res_list"][0]
+    assert len(res["dt_polys"]) > 0
+    assert isinstance(res["rec_texts"], list)
+    assert len(res["rec_texts"]) > 0
+    for text in res["rec_texts"]:
+        assert isinstance(text, str)
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False, "use_doc_unwarping": False},
+        {"use_layout_detection": False},
+        {"layout_det_res": None},
+        {"layout_threshold": 0.5},
+        {"layout_nms": False},
+        {"layout_unclip_ratio": 1.0},
+        {"layout_merge_bboxes_mode": "large"},
+        {"seal_det_limit_side_len": 736},
+        {"seal_det_limit_type": "min"},
+        {"seal_det_thresh": 0.5},
+        {"seal_det_box_thresh": 0.6},
+        {"seal_det_unclip_ratio": 0.5},
+        {"seal_rec_score_thresh": 0.05},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    ocr_engine: SealRecognition,
+    params: dict,
+) -> None:
+    check_wrapper_simple_inference_param_forwarding(
+        monkeypatch,
+        ocr_engine,
+        "paddlex_pipeline",
+        "dummy_path",
+        params,
+    )

tests/pipelines/test_table_recognition_v2.py

@@ -0,0 +1,64 @@
+import pytest
+
+from paddleocr import TableRecognitionPipelineV2
+from ..testing_utils import (
+    TEST_DATA_DIR,
+    check_simple_inference_result,
+    check_wrapper_simple_inference_param_forwarding,
+)
+
+
+@pytest.fixture(scope="module")
+def table_recognition_v2_pipeline():
+    return TableRecognitionPipelineV2()
+
+
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        TEST_DATA_DIR / "table.jpg",
+    ],
+)
+def test_visual_predict(table_recognition_v2_pipeline, image_path):
+    result = table_recognition_v2_pipeline.predict(
+        str(image_path), use_doc_orientation_classify=False, use_doc_unwarping=False
+    )
+
+    check_simple_inference_result(result)
+    res = result[0]
+    assert len(res["table_res_list"]) > 0
+    assert isinstance(res["table_res_list"][0], dict)
+    assert len(res["table_res_list"][0]["cell_box_list"]) > 0
+    assert isinstance(res["table_res_list"][0]["pred_html"], str)
+    assert isinstance(res["table_res_list"][0]["table_ocr_pred"], dict)
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"use_doc_orientation_classify": False},
+        {"use_doc_unwarping": False},
+        {"use_layout_detection": False},
+        {"use_ocr_model": False},
+        {"text_det_limit_side_len": 640, "text_det_limit_type": "min"},
+        {"text_det_thresh": 0.5},
+        {"text_det_box_thresh": 0.3},
+        {"text_det_unclip_ratio": 3.0},
+        {"text_rec_score_thresh": 0.5},
+    ],
+)
+def test_predict_params(
+    monkeypatch,
+    table_recognition_v2_pipeline,
+    params,
+):
+    check_wrapper_simple_inference_param_forwarding(
+        monkeypatch,
+        table_recognition_v2_pipeline,
+        "paddlex_pipeline",
+        "dummy_path",
+        params,
+    )
+
+
+# TODO: Test constructor and other methods