103 lines
2.5 KiB
Python
103 lines
2.5 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 FakeBlock:
|
|
label = "text"
|
|
bbox = [0, 0, 10, 10]
|
|
content = "mock OCR"
|
|
image = None
|
|
|
|
|
|
class FakeResult(dict):
|
|
@property
|
|
def markdown(self):
|
|
return {"markdown_texts": "mock OCR", "markdown_images": {}}
|
|
|
|
@property
|
|
def json(self):
|
|
return {"res": {"input_path": self["input_path"]}}
|
|
|
|
|
|
class FakePipeline:
|
|
def predict(self, path, **kwargs):
|
|
return [
|
|
FakeResult(
|
|
input_path=path,
|
|
width=144,
|
|
height=144,
|
|
layout_det_res={"boxes": [{}]},
|
|
parsing_res_list=[FakeBlock()],
|
|
)
|
|
]
|
|
|
|
|
|
class FakeProvider:
|
|
resolved_device = "cpu"
|
|
|
|
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 metadata(self):
|
|
return {
|
|
"device": "cpu",
|
|
"model_init_seconds": self.model_init_seconds,
|
|
}
|
|
|
|
|
|
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 not result["model_initialized_during_task"]
|
|
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"))
|
|
assert result["text_pages"] == 0
|
|
assert result["ocr_pages"] == 1
|
|
assert result["model_used"]
|
|
assert result["model_initialized_during_task"]
|
|
assert provider.get_calls == 1
|
|
assert manifest["pages"]["1"]["routing_reason"] == "empty_text_layer"
|