127 lines
3.3 KiB
Python
127 lines
3.3 KiB
Python
import json
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
from ocr_app.commands import process_image_file
|
|
from ocr_app.logging_utils import setup_run_logger
|
|
from ocr_app.output import image_output_directory, pdf_output_root
|
|
|
|
|
|
class FakeBlock:
|
|
label = "text"
|
|
bbox = [0, 0, 10, 10]
|
|
content = "hello OCR"
|
|
image = None
|
|
|
|
|
|
class FakeResult(dict):
|
|
@property
|
|
def markdown(self):
|
|
return {"markdown_texts": "hello OCR", "markdown_images": {}}
|
|
|
|
@property
|
|
def json(self):
|
|
return {"res": {"input_path": self["input_path"]}}
|
|
|
|
|
|
class FakePipeline:
|
|
def predict(self, path):
|
|
return [
|
|
FakeResult(
|
|
input_path=path,
|
|
width=20,
|
|
height=10,
|
|
layout_det_res={"boxes": [{}]},
|
|
parsing_res_list=[FakeBlock()],
|
|
)
|
|
]
|
|
|
|
|
|
class FakeProvider:
|
|
class Config:
|
|
device = "cpu"
|
|
|
|
config = Config()
|
|
model_init_seconds = 0.01
|
|
|
|
def get(self):
|
|
return FakePipeline()
|
|
|
|
def synchronize(self):
|
|
pass
|
|
|
|
def metadata(self):
|
|
return {"device": "cpu", "model_init_seconds": self.model_init_seconds}
|
|
|
|
def gpu_memory(self):
|
|
return {}
|
|
|
|
|
|
def make_args(output):
|
|
return Namespace(
|
|
warmup=0,
|
|
rounds=1,
|
|
output=output,
|
|
recursive=False,
|
|
benchmark_json=None,
|
|
no_result=True,
|
|
)
|
|
|
|
|
|
def test_single_image_generates_output_files(tmp_path):
|
|
image = tmp_path / "card.jpg"
|
|
Image.new("RGB", (20, 10), "white").save(image)
|
|
logger = setup_run_logger("test.image.output", tmp_path / "run.log", console=False)
|
|
|
|
result = process_image_file(
|
|
image,
|
|
args=make_args(tmp_path / "outputs"),
|
|
provider=FakeProvider(),
|
|
logger=logger,
|
|
project_root=tmp_path,
|
|
run_warmup=True,
|
|
batch_root=None,
|
|
)
|
|
|
|
output_dir = Path(result.details["output_dir"])
|
|
assert output_dir == tmp_path / "outputs" / "images" / "card_jpg"
|
|
assert (output_dir / "result.md").read_text("utf-8").strip() == "hello OCR"
|
|
assert (output_dir / "result.txt").read_text("utf-8").strip() == "hello OCR"
|
|
data = json.loads((output_dir / "result.json").read_text("utf-8"))
|
|
assert data["res"]["source_type"] == "image_ocr"
|
|
benchmark = json.loads((output_dir / "benchmark.json").read_text("utf-8"))
|
|
assert benchmark["file_total_seconds"] >= benchmark["processing_seconds"]
|
|
|
|
|
|
def test_recursive_output_paths_preserve_relative_directories(tmp_path):
|
|
batch_root = tmp_path / "input"
|
|
image = batch_root / "sub" / "same.png"
|
|
pdf = batch_root / "other" / "same.pdf"
|
|
image.parent.mkdir(parents=True)
|
|
pdf.parent.mkdir(parents=True)
|
|
output = tmp_path / "outputs"
|
|
|
|
assert image_output_directory(
|
|
output,
|
|
image,
|
|
batch_root=batch_root,
|
|
recursive=True,
|
|
) == output / "images" / "sub" / "same_png"
|
|
assert pdf_output_root(
|
|
output,
|
|
pdf,
|
|
batch_root=batch_root,
|
|
recursive=True,
|
|
) == output / "pdfs" / "other"
|
|
|
|
|
|
def test_image_extensions_do_not_collide(tmp_path):
|
|
output = tmp_path / "outputs"
|
|
png = tmp_path / "same.png"
|
|
jpg = tmp_path / "same.jpg"
|
|
assert image_output_directory(output, png, batch_root=None, recursive=False) != image_output_directory(
|
|
output, jpg, batch_root=None, recursive=False
|
|
)
|