116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
"""Output helpers for normalized PP-OCRv6 results."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
from .result_adapter import raw_result_json, result_markdown, result_plain_text
|
|
|
|
|
|
def safe_stem(value: str) -> str:
|
|
cleaned = re.sub(r"[^\w.-]+", "_", value, flags=re.UNICODE).strip("._")
|
|
return cleaned or "result"
|
|
|
|
|
|
def atomic_write_text(path: Path, content: str) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
temporary.write_text(content, encoding="utf-8")
|
|
temporary.replace(path)
|
|
|
|
|
|
def atomic_write_json(path: Path, data: Any) -> None:
|
|
atomic_write_text(path, json.dumps(data, ensure_ascii=False, indent=2))
|
|
|
|
|
|
def image_output_directory(
|
|
output_root: Path,
|
|
image_path: Path,
|
|
*,
|
|
batch_root: Path | None,
|
|
recursive: bool,
|
|
) -> Path:
|
|
base = output_root.expanduser().resolve() / "images"
|
|
if batch_root is not None and recursive:
|
|
base /= image_path.parent.resolve().relative_to(batch_root.resolve())
|
|
suffix = image_path.suffix.lower().lstrip(".") or "image"
|
|
return base / safe_stem(f"{image_path.stem}_{suffix}")
|
|
|
|
|
|
def pdf_output_root(
|
|
output_root: Path,
|
|
pdf_path: Path,
|
|
*,
|
|
batch_root: Path | None,
|
|
recursive: bool,
|
|
) -> Path:
|
|
base = output_root.expanduser().resolve() / "pdfs"
|
|
if batch_root is not None and recursive:
|
|
base /= pdf_path.parent.resolve().relative_to(batch_root.resolve())
|
|
return base
|
|
|
|
|
|
def _save_visualization(result: Any, path: Path) -> bool:
|
|
try:
|
|
images = result.img
|
|
except Exception:
|
|
return False
|
|
if isinstance(images, dict):
|
|
image = images.get("ocr_res_img")
|
|
if image is None:
|
|
image = next(iter(images.values()), None)
|
|
else:
|
|
image = images
|
|
if image is None:
|
|
return False
|
|
if not isinstance(image, Image.Image):
|
|
try:
|
|
image = Image.fromarray(image)
|
|
except Exception:
|
|
return False
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
image.convert("RGB").save(temporary, format="JPEG", quality=92)
|
|
temporary.replace(path)
|
|
return True
|
|
|
|
|
|
def save_image_ocr_outputs(
|
|
result: Any,
|
|
normalized: dict[str, Any],
|
|
output_dir: Path,
|
|
*,
|
|
input_path: Path,
|
|
benchmark: dict[str, Any],
|
|
save_raw_result: bool,
|
|
save_visualization: bool,
|
|
) -> dict[str, str]:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
paths: dict[str, str] = {
|
|
"output_dir": str(output_dir),
|
|
"markdown": str(output_dir / "result.md"),
|
|
"text": str(output_dir / "result.txt"),
|
|
"json": str(output_dir / "result.json"),
|
|
"benchmark": str(output_dir / "benchmark.json"),
|
|
}
|
|
atomic_write_text(Path(paths["markdown"]), result_markdown(normalized, title=input_path.name))
|
|
text = result_plain_text(normalized)
|
|
atomic_write_text(Path(paths["text"]), text.rstrip() + ("\n" if text else ""))
|
|
atomic_write_json(Path(paths["json"]), normalized)
|
|
atomic_write_json(Path(paths["benchmark"]), benchmark)
|
|
|
|
if save_raw_result:
|
|
raw_path = output_dir / "raw-result.json"
|
|
atomic_write_json(raw_path, raw_result_json(result))
|
|
paths["raw_json"] = str(raw_path)
|
|
if save_visualization:
|
|
visualization = output_dir / "visualization.jpg"
|
|
if _save_visualization(result, visualization):
|
|
paths["visualization"] = str(visualization)
|
|
return paths
|