117 lines
3.5 KiB
Python
117 lines
3.5 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 FakeResult(dict):
|
|
@property
|
|
def json(self):
|
|
return dict(self)
|
|
|
|
|
|
class FakePipeline:
|
|
def predict(self, path, **kwargs):
|
|
return [
|
|
FakeResult(
|
|
input_path=path,
|
|
page_index=0,
|
|
rec_texts=["hello OCR"],
|
|
rec_scores=[0.99],
|
|
rec_polys=[[[0, 0], [10, 0], [10, 5], [0, 5]]],
|
|
rec_boxes=[[0, 0, 10, 5]],
|
|
)
|
|
]
|
|
|
|
|
|
class FakeProvider:
|
|
class Config:
|
|
device = "cpu"
|
|
lang = "ch"
|
|
model_size = "medium"
|
|
text_detection_model_name = "PP-OCRv6_medium_det"
|
|
text_recognition_model_name = "PP-OCRv6_medium_rec"
|
|
|
|
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,
|
|
save_raw_result=False,
|
|
save_visualization=False,
|
|
text_det_limit_side_len=None,
|
|
text_det_limit_type=None,
|
|
text_det_thresh=None,
|
|
text_det_box_thresh=None,
|
|
text_det_unclip_ratio=None,
|
|
text_rec_score_thresh=0.0,
|
|
return_word_box=False,
|
|
)
|
|
|
|
|
|
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 "hello OCR" in (output_dir / "result.md").read_text("utf-8")
|
|
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["source_type"] == "image_ocr"
|
|
assert data["lines"][0]["score"] == 0.99
|
|
benchmark = json.loads((output_dir / "benchmark.json").read_text("utf-8"))
|
|
assert benchmark["file_total_seconds"] >= benchmark["export_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"
|
|
assert image_output_directory(output, tmp_path / "same.png", batch_root=None, recursive=False) != image_output_directory(output, tmp_path / "same.jpg", batch_root=None, recursive=False)
|