398 lines
18 KiB
Python
398 lines
18 KiB
Python
"""Hybrid PDF processing with PP-OCRv6 fallback for scanned pages."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import re
|
|
import shutil
|
|
import time
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
import pypdfium2 as pdfium
|
|
from PIL import Image
|
|
|
|
from .output import atomic_write_json, atomic_write_text, safe_stem
|
|
from .pdf_text import TextLayerPolicy, extract_page_text
|
|
from .result_adapter import adapt_ocr_result, result_markdown
|
|
|
|
MANIFEST_VERSION = 1
|
|
PAGE_SPEC_PATTERN = re.compile(r"^(\d+)(?:-(\d*)?)?$")
|
|
PDF_MODES = {"hybrid", "text", "ocr"}
|
|
|
|
|
|
def now_iso() -> str:
|
|
return datetime.now().astimezone().isoformat()
|
|
|
|
|
|
def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as file:
|
|
while chunk := file.read(chunk_size):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def parse_page_spec(spec: str | None, page_count: int) -> list[int]:
|
|
if page_count < 1:
|
|
return []
|
|
if spec is None or not spec.strip():
|
|
return list(range(page_count))
|
|
selected: set[int] = set()
|
|
for raw_part in spec.split(","):
|
|
part = raw_part.strip()
|
|
match = PAGE_SPEC_PATTERN.fullmatch(part)
|
|
if not match:
|
|
raise ValueError(f"无效页码范围: {part!r}，示例: 1-5,8,10-")
|
|
start = int(match.group(1))
|
|
end_text = match.group(2)
|
|
end = start if "-" not in part else int(end_text) if end_text else page_count
|
|
if start < 1 or end < 1:
|
|
raise ValueError("PDF 页码从 1 开始")
|
|
if start > end:
|
|
raise ValueError(f"页码起始值不能大于结束值: {part}")
|
|
if start > page_count or end > page_count:
|
|
raise ValueError(f"页码范围 {part} 超出 PDF 总页数 {page_count}")
|
|
selected.update(range(start - 1, end))
|
|
return sorted(selected)
|
|
|
|
|
|
def render_page(document: Any, page_index: int, dpi: int) -> Image.Image:
|
|
page = document.get_page(page_index)
|
|
bitmap = None
|
|
try:
|
|
bitmap = page.render(scale=dpi / 72.0)
|
|
return bitmap.to_pil().convert("RGB").copy()
|
|
finally:
|
|
if bitmap is not None:
|
|
bitmap.close()
|
|
page.close()
|
|
|
|
|
|
def save_png_atomic(image: Image.Image, path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
image.save(temporary, format="PNG")
|
|
temporary.replace(path)
|
|
|
|
|
|
def _page_paths(document_dir: Path, page_number: int) -> tuple[Path, Path]:
|
|
stem = f"page-{page_number:04d}"
|
|
return document_dir / "pages" / f"{stem}.md", document_dir / "pages" / f"{stem}.json"
|
|
|
|
|
|
def _page_is_complete(document_dir: Path, manifest: dict[str, Any], page_number: int) -> bool:
|
|
record = manifest.get("pages", {}).get(str(page_number), {})
|
|
markdown_path, json_path = _page_paths(document_dir, page_number)
|
|
return record.get("status") == "completed" and markdown_path.is_file() and json_path.is_file()
|
|
|
|
|
|
def rebuild_combined_outputs(document_dir: Path, manifest: dict[str, Any]) -> None:
|
|
markdown_parts = [f"# {manifest['document_name']}"]
|
|
page_results = []
|
|
for page_number in manifest.get("selected_pages", []):
|
|
record = manifest.get("pages", {}).get(str(page_number), {})
|
|
markdown_path, json_path = _page_paths(document_dir, page_number)
|
|
if record.get("status") == "completed" and markdown_path.is_file() and json_path.is_file():
|
|
source = record.get("source_type", "unknown")
|
|
text = markdown_path.read_text(encoding="utf-8").strip()
|
|
markdown_parts.append(f"\n\n---\n\n## Page {page_number} ({source})\n\n{text}")
|
|
page_results.append({"page_number": page_number, "source_type": source, "metrics": record, "result": json.loads(json_path.read_text(encoding="utf-8"))})
|
|
elif record.get("status") == "failed":
|
|
markdown_parts.append(f"\n\n---\n\n## Page {page_number}\n\n> Failed: {record.get('error')}")
|
|
atomic_write_text(document_dir / "document.md", "".join(markdown_parts).rstrip() + "\n")
|
|
atomic_write_json(document_dir / "document.json", {"manifest": manifest, "page_results": page_results})
|
|
|
|
|
|
def validate_pdf_request(pdf_path: Path, output_root: Path, *, resume: bool, overwrite: bool) -> tuple[Path, Path]:
|
|
pdf_path = pdf_path.expanduser().resolve()
|
|
output_root = output_root.expanduser().resolve()
|
|
if not pdf_path.is_file():
|
|
raise FileNotFoundError(f"PDF 不存在: {pdf_path}")
|
|
if pdf_path.suffix.lower() != ".pdf":
|
|
raise ValueError(f"输入文件不是 PDF: {pdf_path}")
|
|
if resume and overwrite:
|
|
raise ValueError("--resume 和 --overwrite 不能同时使用")
|
|
document_dir = output_root / safe_stem(pdf_path.stem)
|
|
if resume and not (document_dir / "manifest.json").is_file():
|
|
raise FileNotFoundError(f"无法续传，缺少 {document_dir / 'manifest.json'}")
|
|
if document_dir.exists() and any(document_dir.iterdir()) and not (resume or overwrite):
|
|
raise FileExistsError(f"输出目录已存在: {document_dir}；请使用 --resume 或 --overwrite")
|
|
return pdf_path, output_root
|
|
|
|
|
|
def preflight_pdf(*, pdf_path: Path, output_root: Path, pages: str | None, dpi: int, password: str | None, resume: bool, overwrite: bool) -> dict[str, Any]:
|
|
pdf_path, output_root = validate_pdf_request(pdf_path, output_root, resume=resume, overwrite=overwrite)
|
|
if dpi < 72 or dpi > 600:
|
|
raise ValueError("--dpi 必须在 72 到 600 之间")
|
|
document = pdfium.PdfDocument(str(pdf_path), password=password)
|
|
try:
|
|
page_count = len(document)
|
|
selected = parse_page_spec(pages, page_count)
|
|
finally:
|
|
document.close()
|
|
return {"pdf_path": pdf_path, "output_root": output_root, "document_dir": output_root / safe_stem(pdf_path.stem), "page_count": page_count, "selected_pages": [index + 1 for index in selected]}
|
|
|
|
|
|
def _prepare_manifest(
|
|
*,
|
|
pdf_path: Path,
|
|
document_dir: Path,
|
|
page_count: int,
|
|
selected_pages: Iterable[int],
|
|
dpi: int,
|
|
mode: str,
|
|
policy: TextLayerPolicy,
|
|
resume: bool,
|
|
overwrite: bool,
|
|
model_config: dict[str, Any],
|
|
run_metadata: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
manifest_path = document_dir / "manifest.json"
|
|
digest = sha256_file(pdf_path)
|
|
selected = [index + 1 for index in selected_pages]
|
|
if overwrite and document_dir.exists():
|
|
shutil.rmtree(document_dir)
|
|
if resume:
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
if manifest.get("manifest_version") != MANIFEST_VERSION or manifest.get("pipeline") != "PP-OCRv6":
|
|
raise ValueError("manifest 与当前 PP-OCRv6 项目不兼容，请使用 --overwrite")
|
|
if manifest.get("input", {}).get("sha256") != digest:
|
|
raise ValueError("PDF 内容已变化，请使用 --overwrite")
|
|
if manifest.get("render", {}).get("dpi") != dpi or manifest.get("mode") != mode:
|
|
raise ValueError("DPI 或模式与原任务不一致，请使用原参数或 --overwrite")
|
|
if manifest.get("text_layer_policy") != policy.__dict__:
|
|
raise ValueError("文本层阈值与原任务不一致，请使用原参数或 --overwrite")
|
|
if manifest.get("model_config") != model_config:
|
|
raise ValueError("PP-OCRv6 模型配置与原任务不一致，请使用原参数或 --overwrite")
|
|
manifest["selected_pages"] = sorted(set(manifest.get("selected_pages", [])) | set(selected))
|
|
manifest["run_metadata"] = run_metadata
|
|
manifest["status"] = "running"
|
|
manifest["updated_at"] = now_iso()
|
|
else:
|
|
document_dir.mkdir(parents=True, exist_ok=True)
|
|
manifest = {
|
|
"manifest_version": MANIFEST_VERSION,
|
|
"pipeline": "PP-OCRv6",
|
|
"document_name": pdf_path.stem,
|
|
"input": {"path": str(pdf_path), "sha256": digest, "size_bytes": pdf_path.stat().st_size},
|
|
"page_count": page_count,
|
|
"selected_pages": selected,
|
|
"mode": mode,
|
|
"model_config": model_config,
|
|
"text_layer_policy": policy.__dict__,
|
|
"render": {"dpi": dpi, "format": "png"},
|
|
"run_metadata": run_metadata,
|
|
"status": "running",
|
|
"created_at": now_iso(),
|
|
"updated_at": now_iso(),
|
|
"pages": {},
|
|
}
|
|
atomic_write_json(manifest_path, manifest)
|
|
return manifest
|
|
|
|
|
|
def process_pdf(
|
|
*,
|
|
provider: Any,
|
|
pdf_path: Path,
|
|
output_root: Path,
|
|
mode: str = "hybrid",
|
|
text_policy: TextLayerPolicy | None = None,
|
|
pages: str | None = None,
|
|
dpi: int = 144,
|
|
password: str | None = None,
|
|
resume: bool = False,
|
|
overwrite: bool = False,
|
|
keep_rendered: bool = False,
|
|
fail_fast: bool = False,
|
|
predict_kwargs: dict[str, Any] | None = None,
|
|
logger: logging.Logger | None = None,
|
|
) -> dict[str, Any]:
|
|
if mode not in PDF_MODES:
|
|
raise ValueError(f"不支持的 PDF 模式: {mode}")
|
|
task_started = time.perf_counter()
|
|
model_init_before = provider.model_init_seconds
|
|
logger = logger or logging.getLogger(__name__)
|
|
text_policy = text_policy or TextLayerPolicy()
|
|
predict_kwargs = predict_kwargs or {}
|
|
pdf_path, output_root = validate_pdf_request(pdf_path, output_root, resume=resume, overwrite=overwrite)
|
|
if dpi < 72 or dpi > 600:
|
|
raise ValueError("--dpi 必须在 72 到 600 之间")
|
|
|
|
document_dir = output_root / safe_stem(pdf_path.stem)
|
|
manifest_path = document_dir / "manifest.json"
|
|
cache_dir = document_dir / ".render-cache"
|
|
document = pdfium.PdfDocument(str(pdf_path), password=password)
|
|
try:
|
|
page_count = len(document)
|
|
selected_indexes = parse_page_spec(pages, page_count)
|
|
manifest = _prepare_manifest(
|
|
pdf_path=pdf_path,
|
|
document_dir=document_dir,
|
|
page_count=page_count,
|
|
selected_pages=selected_indexes,
|
|
dpi=dpi,
|
|
mode=mode,
|
|
policy=text_policy,
|
|
resume=resume,
|
|
overwrite=overwrite,
|
|
model_config=provider.model_config(),
|
|
run_metadata={"device": provider.resolved_device},
|
|
)
|
|
selected_indexes = [number - 1 for number in manifest["selected_pages"]]
|
|
completed_before = sum(_page_is_complete(document_dir, manifest, index + 1) for index in selected_indexes)
|
|
pending = [index for index in selected_indexes if not _page_is_complete(document_dir, manifest, index + 1)]
|
|
|
|
for position, page_index in enumerate(pending, 1):
|
|
page_number = page_index + 1
|
|
started = time.perf_counter()
|
|
text_extract_seconds = render_seconds = ocr_seconds = export_seconds = 0.0
|
|
source_type = "unknown"
|
|
assessment_dict: dict[str, Any] = {}
|
|
render_path = (document_dir / "rendered" if keep_rendered else cache_dir) / f"page-{page_number:04d}.png"
|
|
try:
|
|
text_started = time.perf_counter()
|
|
page = document.get_page(page_index)
|
|
try:
|
|
extracted_text, assessment = extract_page_text(page, text_policy)
|
|
width_points, height_points = page.get_size()
|
|
finally:
|
|
page.close()
|
|
text_extract_seconds = time.perf_counter() - text_started
|
|
assessment_dict = assessment.to_dict()
|
|
use_text = mode == "text" or (mode == "hybrid" and assessment.usable)
|
|
source_type = "text" if use_text else "ocr"
|
|
markdown_path, json_path = _page_paths(document_dir, page_number)
|
|
|
|
if source_type == "text":
|
|
markdown_text = extracted_text.rstrip() + ("\n" if extracted_text else "")
|
|
payload = {
|
|
"schema_version": 1,
|
|
"source_type": "text",
|
|
"input_path": str(pdf_path),
|
|
"page_index": page_index,
|
|
"page_number": page_number,
|
|
"page_count": page_count,
|
|
"text": extracted_text,
|
|
"text_layer": assessment_dict,
|
|
"page_size_points": {"width": width_points, "height": height_points},
|
|
}
|
|
detected_lines = sum(bool(line.strip()) for line in extracted_text.splitlines())
|
|
mean_score = None
|
|
else:
|
|
rendered = time.perf_counter()
|
|
image = render_page(document, page_index, dpi)
|
|
try:
|
|
save_png_atomic(image, render_path)
|
|
finally:
|
|
image.close()
|
|
render_seconds = time.perf_counter() - rendered
|
|
pipeline = provider.get()
|
|
provider.synchronize()
|
|
ocr_started = time.perf_counter()
|
|
results = pipeline.predict(str(render_path), **predict_kwargs)
|
|
provider.synchronize()
|
|
ocr_seconds = time.perf_counter() - ocr_started
|
|
if not results:
|
|
raise RuntimeError("PP-OCRv6 未返回 PDF 页面结果")
|
|
payload = adapt_ocr_result(
|
|
results[0],
|
|
input_path=pdf_path,
|
|
source_type="pdf_ocr",
|
|
language=provider.config.lang,
|
|
detection_model=provider.config.text_detection_model_name,
|
|
recognition_model=provider.config.text_recognition_model_name,
|
|
model_size=provider.config.model_size,
|
|
page_index=page_index,
|
|
page_number=page_number,
|
|
)
|
|
payload["page_count"] = page_count
|
|
payload["ocr_reason"] = assessment.reason if mode == "hybrid" else "forced_ocr_mode"
|
|
payload["text_layer"] = assessment_dict
|
|
payload["render_dpi"] = dpi
|
|
markdown_text = result_markdown(payload)
|
|
detected_lines = payload["summary"]["detected_lines"]
|
|
mean_score = payload["summary"]["mean_score"]
|
|
|
|
export_started = time.perf_counter()
|
|
atomic_write_text(markdown_path, markdown_text)
|
|
atomic_write_json(json_path, payload)
|
|
export_seconds = time.perf_counter() - export_started
|
|
total_seconds = time.perf_counter() - started
|
|
manifest["pages"][str(page_number)] = {
|
|
"status": "completed",
|
|
"page_number": page_number,
|
|
"source_type": source_type,
|
|
"routing_reason": assessment.reason if mode == "hybrid" else f"forced_{source_type}_mode",
|
|
"text_layer": assessment_dict,
|
|
"text_extract_seconds": round(text_extract_seconds, 3),
|
|
"render_seconds": round(render_seconds, 3),
|
|
"ocr_seconds": round(ocr_seconds, 3),
|
|
"export_seconds": round(export_seconds, 3),
|
|
"total_seconds": round(total_seconds, 3),
|
|
"detected_lines": detected_lines,
|
|
"mean_score": mean_score,
|
|
"completed_at": now_iso(),
|
|
}
|
|
except KeyboardInterrupt:
|
|
manifest["status"] = "interrupted"
|
|
atomic_write_json(manifest_path, manifest)
|
|
rebuild_combined_outputs(document_dir, manifest)
|
|
raise
|
|
except Exception as exc:
|
|
manifest["pages"][str(page_number)] = {
|
|
"status": "failed",
|
|
"page_number": page_number,
|
|
"source_type": source_type,
|
|
"text_layer": assessment_dict,
|
|
"total_seconds": round(time.perf_counter() - started, 3),
|
|
"error": f"{type(exc).__name__}: {exc}",
|
|
"failed_at": now_iso(),
|
|
}
|
|
logger.exception("PAGE_FAILED page=%d source=%s", page_number, source_type)
|
|
if fail_fast:
|
|
raise
|
|
finally:
|
|
if not keep_rendered and render_path.is_file():
|
|
render_path.unlink()
|
|
|
|
manifest["run_metadata"] = provider.metadata()
|
|
manifest["updated_at"] = now_iso()
|
|
atomic_write_json(manifest_path, manifest)
|
|
rebuild_combined_outputs(document_dir, manifest)
|
|
logger.info("PAGE_FINISHED page=%d source=%s progress=%d/%d", page_number, source_type, position, len(pending))
|
|
|
|
if cache_dir.exists():
|
|
shutil.rmtree(cache_dir, ignore_errors=True)
|
|
records = [manifest.get("pages", {}).get(str(index + 1), {}) for index in selected_indexes]
|
|
failed_pages = [record.get("page_number") for record in records if record.get("status") == "failed"]
|
|
completed = [record for record in records if record.get("status") == "completed"]
|
|
text_pages = sum(record.get("source_type") == "text" for record in completed)
|
|
ocr_pages = sum(record.get("source_type") == "ocr" for record in completed)
|
|
manifest["status"] = "completed_with_errors" if failed_pages else "completed"
|
|
manifest["run_metadata"] = provider.metadata()
|
|
manifest["summary"] = {
|
|
"selected_pages": len(selected_indexes),
|
|
"completed_pages": len(completed),
|
|
"completed_before_resume": completed_before,
|
|
"text_pages": text_pages,
|
|
"ocr_pages": ocr_pages,
|
|
"failed_pages": failed_pages,
|
|
"model_used": ocr_pages > 0,
|
|
"model_initialized_during_task": model_init_before == 0 and provider.model_init_seconds > 0,
|
|
"timing": {
|
|
"model_init_seconds": round(provider.model_init_seconds, 3),
|
|
"task_total_seconds": round(time.perf_counter() - task_started, 3),
|
|
},
|
|
}
|
|
manifest["updated_at"] = now_iso()
|
|
atomic_write_json(manifest_path, manifest)
|
|
rebuild_combined_outputs(document_dir, manifest)
|
|
return {"document_dir": str(document_dir), "manifest_path": str(manifest_path), "status": manifest["status"], **manifest["summary"]}
|
|
finally:
|
|
document.close()
|