90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
from ocr_app.pdf import process_pdf
|
|
from ocr_app.pdf_text import TextLayerPolicy
|
|
|
|
|
|
class FakeResult(dict):
|
|
pass
|
|
|
|
|
|
class FakePipeline:
|
|
def predict(self, path, **kwargs):
|
|
return [
|
|
FakeResult(
|
|
input_path=path,
|
|
page_index=0,
|
|
rec_texts=["mock OCR"],
|
|
rec_scores=[0.97],
|
|
rec_polys=[[[0, 0], [20, 0], [20, 10], [0, 10]]],
|
|
rec_boxes=[[0, 0, 20, 10]],
|
|
)
|
|
]
|
|
|
|
|
|
class FakeProvider:
|
|
resolved_device = "cpu"
|
|
|
|
class Config:
|
|
lang = "ch"
|
|
model_size = "medium"
|
|
text_detection_model_name = "PP-OCRv6_medium_det"
|
|
text_recognition_model_name = "PP-OCRv6_medium_rec"
|
|
|
|
config = Config()
|
|
|
|
def __init__(self):
|
|
self.model_init_seconds = 0.0
|
|
self.get_calls = 0
|
|
|
|
def get(self):
|
|
self.get_calls += 1
|
|
self.model_init_seconds = 0.01
|
|
return FakePipeline()
|
|
|
|
def synchronize(self):
|
|
pass
|
|
|
|
def model_config(self):
|
|
return {
|
|
"ocr_version": "PP-OCRv6",
|
|
"model_size": "medium",
|
|
"language": "ch",
|
|
"detection_model": "PP-OCRv6_medium_det",
|
|
"recognition_model": "PP-OCRv6_medium_rec",
|
|
}
|
|
|
|
def metadata(self):
|
|
return {"device": "cpu", "model_init_seconds": self.model_init_seconds, "model_config": self.model_config()}
|
|
|
|
|
|
def test_electronic_pdf_does_not_load_model(tmp_path):
|
|
project_root = Path(__file__).resolve().parent.parent
|
|
source = next((project_root / "data" / "documents").glob("*.pdf"))
|
|
provider = FakeProvider()
|
|
|
|
result = process_pdf(provider=provider, pdf_path=source, output_root=tmp_path, mode="hybrid", pages="1", text_policy=TextLayerPolicy())
|
|
|
|
assert result["text_pages"] == 1
|
|
assert result["ocr_pages"] == 0
|
|
assert not result["model_used"]
|
|
assert provider.get_calls == 0
|
|
|
|
|
|
def test_scanned_pdf_falls_back_to_ocr(tmp_path):
|
|
source = tmp_path / "scan.pdf"
|
|
Image.new("RGB", (72, 72), "white").save(source)
|
|
provider = FakeProvider()
|
|
|
|
result = process_pdf(provider=provider, pdf_path=source, output_root=tmp_path / "output", mode="hybrid", text_policy=TextLayerPolicy())
|
|
|
|
manifest = json.loads(Path(result["manifest_path"]).read_text(encoding="utf-8"))
|
|
page_json = json.loads((Path(result["document_dir"]) / "pages" / "page-0001.json").read_text("utf-8"))
|
|
assert result["ocr_pages"] == 1
|
|
assert provider.get_calls == 1
|
|
assert manifest["pages"]["1"]["routing_reason"] == "empty_text_layer"
|
|
assert page_json["lines"][0]["text"] == "mock OCR"
|