PaddleOCR-VL-1.6_Demo/pdf_ocr_core.py

659 lines
26 KiB
Python

"""Shared PDF rendering, OCR orchestration, resume, and export logic."""
from __future__ import annotations
import hashlib
import json
import logging
import os
import re
import shutil
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Iterable
import pypdfium2 as pdfium
from PIL import Image
MANIFEST_VERSION = 1
PAGE_SPEC_PATTERN = re.compile(r"^(\d+)(?:-(\d*)?)?$")
def now_iso() -> str:
return datetime.now().astimezone().isoformat()
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 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 safe_stem(value: str) -> str:
cleaned = re.sub(r"[^\w.-]+", "_", value, flags=re.UNICODE).strip("._")
return cleaned or "document"
def parse_page_spec(spec: str | None, page_count: int) -> list[int]:
"""Parse one-based ranges such as ``1-5,8,10-`` into zero-based indexes."""
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)
if "-" not in part:
end = start
elif end_text:
end = int(end_text)
else:
end = 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 _save_markdown_image(data: Any, path: Path) -> Path:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(f".{path.name}.tmp")
if isinstance(data, Image.Image):
image = data
else:
try:
import numpy as np
array = np.asarray(data)
if array.ndim == 3 and array.shape[2] == 4:
image = Image.fromarray(array.astype("uint8"), mode="RGBA")
elif array.ndim in (2, 3):
image = Image.fromarray(array.astype("uint8"))
else:
raise TypeError(f"unsupported image array shape: {array.shape}")
except Exception as exc:
raise TypeError(f"无法保存 Markdown 图片 {path.name}: {type(data).__name__}") from exc
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"}:
image_format = "PNG"
path = path.with_suffix(".png")
temporary = path.with_name(f".{path.name}.tmp")
image.save(temporary, format=image_format)
temporary.replace(path)
return path
def _result_markdown(result: Any, document_dir: Path, page_number: int) -> str:
markdown_data = result.markdown
if "res" in markdown_data and isinstance(markdown_data["res"], dict):
markdown_data = markdown_data["res"]
text = str(markdown_data.get("markdown_texts", ""))
markdown_images = markdown_data.get("markdown_images") or {}
page_asset_dir = document_dir / "assets" / f"page-{page_number:04d}"
if page_asset_dir.exists():
shutil.rmtree(page_asset_dir)
for index, (original_path, image_data) in enumerate(markdown_images.items(), start=1):
original = str(original_path).replace("\\", "/")
original_name = Path(original).name or f"image-{index:03d}.png"
asset_name = f"{index:03d}-{safe_stem(Path(original_name).stem)}{Path(original_name).suffix or '.png'}"
target = page_asset_dir / asset_name
target = _save_markdown_image(image_data, target)
page_relative = Path(os.path.relpath(target, document_dir / "pages")).as_posix()
text = text.replace(original, page_relative)
text = text.replace(str(original_path), page_relative)
return text.strip()
def _result_json(result: Any) -> dict[str, Any]:
data = result.json
if not isinstance(data, dict):
raise TypeError(f"OCR JSON 结果类型异常: {type(data).__name__}")
return data
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_json_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():
page_text = markdown_path.read_text(encoding="utf-8")
page_text = page_text.replace("../assets/", "assets/")
markdown_parts.append(f"\n\n---\n\n## Page {page_number}\n\n{page_text.strip()}")
page_json_results.append(
{
"page_number": page_number,
"metrics": record,
"ocr_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> OCR failed: {record.get('error', 'unknown 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_json_results,
},
)
def prepare_manifest(
*,
pdf_path: Path,
document_dir: Path,
page_count: int,
selected_pages: Iterable[int],
dpi: int,
resume: bool,
overwrite: bool,
run_metadata: dict[str, Any],
) -> dict[str, Any]:
manifest_path = document_dir / "manifest.json"
pdf_sha256 = sha256_file(pdf_path)
selected_one_based = [index + 1 for index in selected_pages]
if overwrite and document_dir.exists():
shutil.rmtree(document_dir)
if document_dir.exists() and any(document_dir.iterdir()) and not resume:
raise FileExistsError(
f"输出目录已存在: {document_dir}。请使用 --resume 继续或 --overwrite 重建。"
)
if resume:
if not manifest_path.is_file():
raise FileNotFoundError(f"无法断点续传,缺少 manifest: {manifest_path}")
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
if manifest.get("input", {}).get("sha256") != pdf_sha256:
raise ValueError("PDF 内容已变化,不能使用现有断点;请使用 --overwrite")
if manifest.get("render", {}).get("dpi") != dpi:
raise ValueError("DPI 与现有任务不一致;请使用原 DPI 或 --overwrite")
manifest["selected_pages"] = sorted(
set(manifest.get("selected_pages", [])) | set(selected_one_based)
)
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,
"document_name": pdf_path.stem,
"input": {
"path": str(pdf_path),
"sha256": pdf_sha256,
"size_bytes": pdf_path.stat().st_size,
},
"page_count": page_count,
"selected_pages": selected_one_based,
"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 validate_pdf_request(
pdf_path: Path,
output_root: Path,
*,
resume: bool,
overwrite: bool,
) -> tuple[Path, Path]:
"""Validate cheap input/output conditions before loading the large model."""
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"无法断点续传,缺少 manifest: {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]:
"""Validate PDF access, page ranges, and output state before model loading."""
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 process_pdf(
*,
pipeline: Any,
pdf_path: Path,
output_root: Path,
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,
run_metadata: dict[str, Any] | None = None,
predict_kwargs: dict[str, Any] | None = None,
synchronize: Callable[[], None] | None = None,
logger: logging.Logger | None = None,
) -> dict[str, Any]:
"""Render and OCR a PDF one page at a time."""
task_started = time.perf_counter()
logger = logger or logging.getLogger(__name__)
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 之间")
predict_kwargs = predict_kwargs or {}
run_metadata = run_metadata or {}
document_dir = output_root / safe_stem(pdf_path.stem)
manifest_path = document_dir / "manifest.json"
temporary_render_dir = document_dir / ".render-cache"
pdf_open_started = time.perf_counter()
document = pdfium.PdfDocument(str(pdf_path), password=password)
pdf_open_seconds = time.perf_counter() - pdf_open_started
logger.info(
"PDF_OPENED path=%s seconds=%.3f dpi=%d resume=%s overwrite=%s keep_rendered=%s",
pdf_path,
pdf_open_seconds,
dpi,
resume,
overwrite,
keep_rendered,
)
try:
page_count = len(document)
selected_indexes = parse_page_spec(pages, page_count)
manifest_started = time.perf_counter()
manifest = prepare_manifest(
pdf_path=pdf_path,
document_dir=document_dir,
page_count=page_count,
selected_pages=selected_indexes,
dpi=dpi,
resume=resume,
overwrite=overwrite,
run_metadata=run_metadata,
)
manifest_prepare_seconds = time.perf_counter() - manifest_started
logger.info(
"MANIFEST_READY path=%s seconds=%.3f page_count=%d requested_pages=%d",
manifest_path,
manifest_prepare_seconds,
page_count,
len(selected_indexes),
)
# Resume uses the union stored in the manifest, so newly added ranges and
# previously selected pages remain one coherent document task.
selected_indexes = [page_number - 1 for page_number in manifest["selected_pages"]]
completed_before = sum(
_page_is_complete(document_dir, manifest, index + 1) for index in selected_indexes
)
pending_indexes = [
index
for index in selected_indexes
if not _page_is_complete(document_dir, manifest, index + 1)
]
logger.info(
"TASK_PLAN total_pages=%d selected_pages=%d completed_before=%d pending_pages=%d output=%s",
page_count,
len(selected_indexes),
completed_before,
len(pending_indexes),
document_dir,
)
run_page_times: list[float] = []
for position, page_index in enumerate(pending_indexes, start=1):
page_number = page_index + 1
page_started = time.perf_counter()
render_seconds = 0.0
ocr_seconds = 0.0
export_seconds = 0.0
state_save_seconds = 0.0
render_path = temporary_render_dir / f"page-{page_number:04d}.png"
if keep_rendered:
render_path = document_dir / "rendered" / f"page-{page_number:04d}.png"
logger.info(
"PAGE_START page=%d page_index=%d position=%d/%d",
page_number,
page_index,
position,
len(pending_indexes),
)
try:
render_started = 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() - render_started
logger.info(
"PAGE_RENDERED page=%d seconds=%.3f path=%s",
page_number,
render_seconds,
render_path,
)
if synchronize:
synchronize()
ocr_started = time.perf_counter()
result_list = pipeline.predict(str(render_path), **predict_kwargs)
if synchronize:
synchronize()
ocr_seconds = time.perf_counter() - ocr_started
logger.info("PAGE_OCR_COMPLETED page=%d seconds=%.3f", page_number, ocr_seconds)
if not result_list:
raise RuntimeError("OCR pipeline 未返回结果")
export_started = time.perf_counter()
result = result_list[0]
markdown_text = _result_markdown(result, document_dir, page_number)
result_json = _result_json(result)
json_payload = result_json.get("res", result_json)
if isinstance(json_payload, dict):
json_payload["input_path"] = str(pdf_path)
json_payload["page_index"] = page_index
json_payload["page_number"] = page_number
json_payload["page_count"] = page_count
json_payload["render_dpi"] = dpi
markdown_path, json_path = _page_paths(document_dir, page_number)
atomic_write_text(markdown_path, markdown_text.rstrip() + "\n")
atomic_write_json(json_path, result_json)
export_seconds = time.perf_counter() - export_started
total_seconds = time.perf_counter() - page_started
manifest["pages"][str(page_number)] = {
"status": "completed",
"page_number": page_number,
"render_seconds": round(render_seconds, 3),
"ocr_seconds": round(ocr_seconds, 3),
"export_seconds": round(export_seconds, 3),
"total_seconds": round(total_seconds, 3),
"width": result.get("width"),
"height": result.get("height"),
"layout_boxes": len(result.get("layout_det_res", {}).get("boxes", [])),
"parsed_blocks": len(result.get("parsing_res_list", [])),
"device": run_metadata.get("device"),
"completed_at": now_iso(),
}
run_page_times.append(total_seconds)
logger.info(
"PAGE_RESULT_SAVED page=%d seconds=%.3f markdown=%s json=%s width=%s height=%s layout_boxes=%d parsed_blocks=%d",
page_number,
export_seconds,
markdown_path,
json_path,
result.get("width"),
result.get("height"),
len(result.get("layout_det_res", {}).get("boxes", [])),
len(result.get("parsing_res_list", [])),
)
except KeyboardInterrupt:
manifest["status"] = "interrupted"
manifest["updated_at"] = now_iso()
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
logger.warning(
"TASK_INTERRUPTED page=%d elapsed_seconds=%.3f",
page_number,
time.perf_counter() - task_started,
)
raise
except Exception as exc:
total_seconds = time.perf_counter() - page_started
manifest["pages"][str(page_number)] = {
"status": "failed",
"page_number": page_number,
"render_seconds": round(render_seconds, 3),
"ocr_seconds": round(ocr_seconds, 3),
"export_seconds": round(export_seconds, 3),
"total_seconds": round(total_seconds, 3),
"error": f"{type(exc).__name__}: {exc}",
"failed_at": now_iso(),
}
logger.exception(
"PAGE_FAILED page=%d render_seconds=%.3f ocr_seconds=%.3f export_seconds=%.3f total_seconds=%.3f",
page_number,
render_seconds,
ocr_seconds,
export_seconds,
total_seconds,
)
if fail_fast:
manifest["status"] = "failed"
manifest["updated_at"] = now_iso()
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
raise
finally:
if not keep_rendered and render_path.is_file():
render_path.unlink()
state_save_started = time.perf_counter()
manifest["updated_at"] = now_iso()
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
state_save_seconds = time.perf_counter() - state_save_started
manifest["pages"][str(page_number)]["state_save_seconds"] = round(state_save_seconds, 3)
atomic_write_json(manifest_path, manifest)
processed_now = position
average = sum(run_page_times) / len(run_page_times) if run_page_times else None
remaining = len(pending_indexes) - processed_now
eta = average * remaining if average is not None else None
record = manifest["pages"][str(page_number)]
elapsed_task = time.perf_counter() - task_started
logger.info(
"PAGE_FINISHED page=%d status=%s render_seconds=%.3f ocr_seconds=%.3f export_seconds=%.3f state_save_seconds=%.3f page_total_seconds=%.3f task_elapsed_seconds=%.3f eta_seconds=%s progress=%d/%d",
page_number,
record["status"],
record.get("render_seconds", 0.0),
record.get("ocr_seconds", 0.0),
record.get("export_seconds", 0.0),
state_save_seconds,
record.get("total_seconds", 0.0),
elapsed_task,
f"{eta:.3f}" if eta is not None else "unknown",
processed_now,
len(pending_indexes),
)
if temporary_render_dir.exists():
shutil.rmtree(temporary_render_dir, ignore_errors=True)
selected_records = [
manifest.get("pages", {}).get(str(index + 1), {}) for index in selected_indexes
]
failed_pages = [
record.get("page_number") for record in selected_records if record.get("status") == "failed"
]
completed_pages = sum(record.get("status") == "completed" for record in selected_records)
completed_records = [record for record in selected_records if record.get("status") == "completed"]
render_total = sum(record.get("render_seconds", 0.0) for record in completed_records)
ocr_total = sum(record.get("ocr_seconds", 0.0) for record in completed_records)
export_total = sum(record.get("export_seconds", 0.0) for record in completed_records)
state_save_total = sum(record.get("state_save_seconds", 0.0) for record in selected_records)
page_total = sum(record.get("total_seconds", 0.0) for record in selected_records)
average_ocr = ocr_total / completed_pages if completed_pages else 0.0
average_page = page_total / len(selected_records) if selected_records else 0.0
manifest["status"] = "completed_with_errors" if failed_pages else "completed"
manifest["summary"] = {
"selected_pages": len(selected_indexes),
"completed_pages": completed_pages,
"completed_before_resume": completed_before,
"failed_pages": failed_pages,
"timing": {
"pdf_open_seconds": round(pdf_open_seconds, 3),
"manifest_prepare_seconds": round(manifest_prepare_seconds, 3),
"render_total_seconds": round(render_total, 3),
"ocr_total_seconds": round(ocr_total, 3),
"export_total_seconds": round(export_total, 3),
"state_save_total_seconds": round(state_save_total, 3),
"page_total_seconds": round(page_total, 3),
"average_ocr_seconds": round(average_ocr, 3),
"average_page_seconds": round(average_page, 3),
"finalize_seconds": 0.0,
"task_total_seconds": 0.0,
},
}
finalize_started = time.perf_counter()
manifest["updated_at"] = now_iso()
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
finalize_seconds = time.perf_counter() - finalize_started
task_total = time.perf_counter() - task_started
manifest["summary"]["timing"]["finalize_seconds"] = round(finalize_seconds, 3)
manifest["summary"]["timing"]["task_total_seconds"] = round(task_total, 3)
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
logger.info(
"TASK_COMPLETED status=%s selected_pages=%d completed_pages=%d failed_pages=%s pdf_open_seconds=%.3f manifest_prepare_seconds=%.3f render_total_seconds=%.3f ocr_total_seconds=%.3f export_total_seconds=%.3f state_save_total_seconds=%.3f page_total_seconds=%.3f average_ocr_seconds=%.3f average_page_seconds=%.3f finalize_seconds=%.3f task_total_seconds=%.3f",
manifest["status"],
len(selected_indexes),
completed_pages,
failed_pages,
pdf_open_seconds,
manifest_prepare_seconds,
render_total,
ocr_total,
export_total,
state_save_total,
page_total,
average_ocr,
average_page,
finalize_seconds,
task_total,
)
return {
"document_dir": str(document_dir),
"manifest_path": str(manifest_path),
"status": manifest["status"],
**manifest["summary"],
}
finally:
document.close()