140 lines
4.5 KiB
Python
140 lines
4.5 KiB
Python
"""Shared output helpers for image OCR results and batch manifests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
|
|
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:
|
|
relative_parent = image_path.parent.resolve().relative_to(batch_root.resolve())
|
|
base /= relative_parent
|
|
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:
|
|
relative_parent = pdf_path.parent.resolve().relative_to(batch_root.resolve())
|
|
base /= relative_parent
|
|
return base
|
|
|
|
|
|
def _save_asset(data: Any, path: Path) -> Path:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
if isinstance(data, Image.Image):
|
|
image = data
|
|
else:
|
|
import numpy as np
|
|
|
|
array = np.asarray(data)
|
|
if array.ndim not in (2, 3):
|
|
raise TypeError(f"不支持的图片资源形状: {array.shape}")
|
|
image = Image.fromarray(array.astype("uint8"))
|
|
|
|
image_format = (path.suffix.lstrip(".") or "png").upper()
|
|
if image_format == "JPG":
|
|
image_format = "JPEG"
|
|
if image_format not in {"PNG", "JPEG", "WEBP", "BMP", "TIFF"}:
|
|
path = path.with_suffix(".png")
|
|
image_format = "PNG"
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
image.save(temporary, format=image_format)
|
|
temporary.replace(path)
|
|
return path
|
|
|
|
|
|
def save_image_ocr_outputs(
|
|
result: Any,
|
|
output_dir: Path,
|
|
*,
|
|
input_path: Path,
|
|
benchmark: dict[str, Any],
|
|
) -> dict[str, str]:
|
|
"""Persist Markdown, plain text, Paddle JSON, and benchmark data."""
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
markdown_data = result.markdown
|
|
if "res" in markdown_data and isinstance(markdown_data["res"], dict):
|
|
markdown_data = markdown_data["res"]
|
|
markdown_text = str(markdown_data.get("markdown_texts", "")).strip()
|
|
asset_dir = output_dir / "assets"
|
|
if asset_dir.exists():
|
|
import shutil
|
|
|
|
shutil.rmtree(asset_dir)
|
|
for index, (original_path, image_data) in enumerate(
|
|
(markdown_data.get("markdown_images") or {}).items(),
|
|
start=1,
|
|
):
|
|
original = str(original_path).replace("\\", "/")
|
|
source_name = Path(original).name or f"image-{index:03d}.png"
|
|
target = asset_dir / (
|
|
f"{index:03d}-{safe_stem(Path(source_name).stem)}"
|
|
f"{Path(source_name).suffix or '.png'}"
|
|
)
|
|
target = _save_asset(image_data, target)
|
|
relative = Path(os.path.relpath(target, output_dir)).as_posix()
|
|
markdown_text = markdown_text.replace(original, relative)
|
|
markdown_text = markdown_text.replace(str(original_path), relative)
|
|
|
|
plain_text = "\n\n".join(
|
|
block.content.strip()
|
|
for block in result["parsing_res_list"]
|
|
if block.content.strip()
|
|
)
|
|
result_json = result.json
|
|
payload = result_json.get("res", result_json) if isinstance(result_json, dict) else None
|
|
if isinstance(payload, dict):
|
|
payload["input_path"] = str(input_path)
|
|
payload["source_type"] = "image_ocr"
|
|
|
|
paths = {
|
|
"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"]), markdown_text.rstrip() + "\n")
|
|
atomic_write_text(Path(paths["text"]), plain_text.rstrip() + "\n")
|
|
atomic_write_json(Path(paths["json"]), result_json)
|
|
atomic_write_json(Path(paths["benchmark"]), benchmark)
|
|
return paths
|