"""Shared PDF rendering, OCR orchestration, resume, and export logic.""" from __future__ import annotations import hashlib import json 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 format_duration(seconds: float | None) -> str: if seconds is None: return "unknown" if seconds < 60: return f"{seconds:.1f}s" return f"{seconds / 60:.1f}min" 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, ) -> dict[str, Any]: """Render and OCR a PDF one page at a time.""" 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" 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, resume=resume, overwrite=overwrite, run_metadata=run_metadata, ) # 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) ] print(f"PDF: {pdf_path}") print(f"Pages: {page_count}, selected: {len(selected_indexes)}, pending: {len(pending_indexes)}") print(f"Output: {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 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" 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 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 if not result_list: raise RuntimeError("OCR pipeline 未返回结果") 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) 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), "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) except KeyboardInterrupt: manifest["status"] = "interrupted" manifest["updated_at"] = now_iso() atomic_write_json(manifest_path, manifest) rebuild_combined_outputs(document_dir, manifest) 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), "total_seconds": round(total_seconds, 3), "error": f"{type(exc).__name__}: {exc}", "failed_at": now_iso(), } print(f"[FAILED] Page {page_number}: {exc}") 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() manifest["updated_at"] = now_iso() atomic_write_json(manifest_path, manifest) rebuild_combined_outputs(document_dir, 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)] print( f"[{processed_now}/{len(pending_indexes)}] Page {page_number}: " f"{record['status']}, OCR {format_duration(record.get('ocr_seconds'))}, " f"ETA {format_duration(eta)}" ) 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) 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, } 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()