diff --git a/.gitignore b/.gitignore index 6c2409d..d09d3ce 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,6 @@ wheels/ # Generated benchmark results benchmarks/gpu/*.json !benchmarks/gpu/.gitkeep + +# OCR outputs +outputs/ diff --git a/README.md b/README.md index c4caf74..a6faa41 100644 --- a/README.md +++ b/README.md @@ -9,11 +9,14 @@ ``` ocr-VL1.6/ ├── main.py # CPU 单图 OCR + Benchmark -├── batch_ocr.py # CPU 批量 OCR（系统友好的多进程版本） +├── batch_ocr.py # CPU 批量图片 OCR（系统友好的多进程版本） +├── pdf_ocr.py # CPU PDF OCR（逐页、可恢复） +├── pdf_ocr_core.py # CPU/GPU 共用的 PDF 渲染、恢复和导出逻辑 ├── pyproject.toml # CPU 项目依赖 ├── uv.lock # CPU 锁文件 ├── gpu/ # 独立 GPU 子项目 │ ├── main.py # GPU 单图 Benchmark +│ ├── pdf_ocr.py # GPU PDF OCR（复用公共核心） │ ├── verify_env.py # CUDA 环境与计算验证 │ ├── setup_env.py # 按 CUDA Wheel 类型创建环境 │ ├── pyproject.toml # GPU 独立依赖 @@ -86,6 +89,93 @@ uv run --project gpu python gpu/main.py --warmup 1 --rounds 3 GPU Benchmark JSON 写入 `benchmarks/gpu/`。详细说明见 [`gpu/README.md`](gpu/README.md)。 +## PDF OCR + +PDF 使用 `pypdfium2` 逐页渲染，再将每一页交给 PaddleOCR-VL。默认采用安全的单进程串行模式，页面完成后立即保存，适合 CPU 长时间任务。CPU 默认预留 2 个逻辑核心给系统，可通过 `--threads` 覆盖。 + +### CPU 使用 + +```bash +# 处理整个 PDF，默认 DPI 144 +uv run python pdf_ocr.py documents/sample.pdf + +# 处理指定页：1-5、8、10 到末页 +uv run python pdf_ocr.py documents/sample.pdf --pages "1-5,8,10-" + +# 中断后继续，已完成页不会重复推理 +uv run python pdf_ocr.py documents/sample.pdf --resume + +# 删除已有输出并重新处理 +uv run python pdf_ocr.py documents/sample.pdf --overwrite + +# 保留每页渲染后的 PNG，便于检查输入质量 +uv run python pdf_ocr.py documents/sample.pdf --keep-rendered + +# 手动设置 CPU 线程数；长任务建议保留 1～2 个核心给系统 +uv run python pdf_ocr.py documents/sample.pdf --threads 18 +``` + +### GPU 使用 + +GPU 入口与 CPU 入口使用相同的 PDF 核心逻辑，但必须在 `gpu/.venv` 和 NVIDIA CUDA GPU 上运行： + +```bash +uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf \ + --device-id 0 \ + --pages "1-10" \ + --dpi 144 +``` + +当前机器无 NVIDIA 独立显卡，因此 GPU PDF 入口仅完成静态检查，尚未实机验证。 + +### 页码语法 + +| 参数 | 含义 | +|------|------| +| `1` | 仅第 1 页 | +| `1-5` | 第 1～5 页 | +| `10-` | 第 10 页到最后一页 | +| `1-5,8,10-` | 多个页码范围组合 | + +用户页码从 1 开始；内部 manifest 使用同样的一基页码记录。 + +### 输出结构 + +```text +outputs/ +└── sample/ + ├── manifest.json # 任务配置、页状态、耗时和错误 + ├── document.md # 合并后的 Markdown + ├── document.json # 合并后的 JSON + ├── pages/ + │ ├── page-0001.md + │ ├── page-0001.json + │ └── ... + ├── assets/ # 表格、图片等 Markdown 资源 + └── rendered/ # 仅使用 --keep-rendered 时保留 +``` + +默认不保留中间渲染 PNG；OCR 完成后会删除临时图。每页 JSON 会将 `input_path` 恢复为原 PDF 路径，并记录 `page_index`、`page_number`、`page_count` 和 `render_dpi`。 + +### 恢复与错误处理 + +- 输出目录已存在时，必须显式使用 `--resume` 或 `--overwrite` +- `--resume` 会校验 PDF SHA-256 和 DPI，防止接续到错误任务 +- 单页失败默认写入 manifest 并继续后续页面 +- `--fail-fast` 可在第一页失败后立即停止 +- `Ctrl+C` 会保存当前 manifest；下次使用 `--resume` 继续 +- 逐页文件和 manifest 使用临时文件替换，降低中途退出造成文件损坏的概率 + +### DPI 建议 + +| 文档类型 | 建议 DPI | +|----------|---------:| +| 普通打印文字 | 120～144 | +| 小字号文档 | 150～200 | +| 手写或低质量扫描件 | 200～250 | + +CPU 当前单图实测约 162 秒。长 PDF 总时间可粗略按 `待处理页数 × 单页耗时` 估算，因此建议先用 `--pages "1"` 测试效果和耗时，再扩大页码范围。DPI 越高通常越慢，不建议默认使用 300 DPI。 + ## 工作原理 `PaddleOCRVL` pipeline 分两阶段： diff --git a/gpu/README.md b/gpu/README.md index 2e5380a..b60a154 100644 --- a/gpu/README.md +++ b/gpu/README.md @@ -96,9 +96,35 @@ Benchmark 会记录： benchmarks/gpu/gpu-benchmark-YYYYMMDD-HHMMSS.json ``` +## PDF OCR + +GPU PDF 入口复用仓库根目录 `pdf_ocr_core.py`，按页渲染、逐页保存并支持断点续传： + +```bash +uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf \ + --device-id 0 \ + --pages "1-10" \ + --dpi 144 +``` + +常用选项： + +```bash +# 中断后继续 +uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf --resume + +# 删除现有输出后重跑 +uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf --overwrite + +# 保留 PDF 页面的渲染 PNG +uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf --keep-rendered +``` + +无 CUDA 时脚本会立即退出，不会自动回落到 CPU。当前开发机器没有 NVIDIA GPU，因此此入口尚未完成 GPU 实机验证。 + ## 当前范围 -当前只实现单 GPU、单图 Benchmark。暂未实现 GPU 多进程批处理，原因是： +当前实现单 GPU、单图 Benchmark 和单 GPU PDF 逐页 OCR。暂未实现 GPU 多进程批处理，原因是： - 同一 GPU 上启动多个模型实例会重复占用显存 - 多进程通常不会线性提升单卡吞吐 diff --git a/gpu/pdf_ocr.py b/gpu/pdf_ocr.py new file mode 100644 index 0000000..a5a04ae --- /dev/null +++ b/gpu/pdf_ocr.py @@ -0,0 +1,153 @@ +"""GPU entry point for page-by-page PaddleOCR-VL PDF recognition.""" + +from __future__ import annotations + +import argparse +import platform +import sys +import time +from pathlib import Path + +GPU_DIR = Path(__file__).resolve().parent +PROJECT_ROOT = GPU_DIR.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from pdf_ocr_core import preflight_pdf, process_pdf + +DEFAULT_OUTPUT = PROJECT_ROOT / "outputs" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="PaddleOCR-VL-1.6 GPU PDF OCR（逐页、可恢复）", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("pdf", type=Path, help="输入 PDF 文件") + parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="输出根目录") + parser.add_argument("--pages", help="一页或多个页码范围，例如 1-5,8,10-") + parser.add_argument("--dpi", type=int, default=144, help="PDF 页面渲染 DPI") + parser.add_argument("--password", help="加密 PDF 密码") + parser.add_argument("--device-id", type=int, default=0, help="CUDA GPU 编号") + parser.add_argument("--resume", action="store_true", help="跳过已完成页，继续现有任务") + parser.add_argument("--overwrite", action="store_true", help="删除已有输出并重新处理") + parser.add_argument("--keep-rendered", action="store_true", help="保留逐页渲染 PNG") + parser.add_argument("--fail-fast", action="store_true", help="任一页失败后立即停止") + parser.add_argument("--max-new-tokens", type=int, default=None, help="限制每个文本块最大生成 token") + parser.add_argument("--min-pixels", type=int, default=None, help="VLM 最小输入像素参数") + parser.add_argument("--max-pixels", type=int, default=None, help="VLM 最大输入像素参数") + return parser.parse_args() + + +def configure_cuda(device_id: int): + try: + import paddle + except ImportError as exc: + raise RuntimeError("未安装 GPU 子项目依赖，请先运行 gpu/setup_env.py。") from exc + + if not paddle.is_compiled_with_cuda(): + raise RuntimeError("当前 PaddlePaddle 不是 CUDA 构建；本程序不会回退到 CPU。") + + try: + device_count = paddle.device.cuda.device_count() + except Exception as exc: + raise RuntimeError(f"无法查询 CUDA 设备: {exc}") from exc + + if device_count < 1: + raise RuntimeError("未检测到 NVIDIA CUDA GPU；本程序不会回退到 CPU。") + if device_id < 0 or device_id >= device_count: + raise RuntimeError(f"GPU {device_id} 不存在，当前检测到 {device_count} 个设备。") + + device = f"gpu:{device_id}" + paddle.set_device(device) + paddle.device.cuda.synchronize(device_id) + try: + name = paddle.device.cuda.get_device_name(device_id) + except Exception: + name = "unknown" + return paddle, device, name + + +def main() -> int: + args = parse_args() + try: + preflight = preflight_pdf( + pdf_path=args.pdf, + output_root=args.output, + pages=args.pages, + dpi=args.dpi, + password=args.password, + resume=args.resume, + overwrite=args.overwrite, + ) + paddle, device, device_name = configure_cuda(args.device_id) + except Exception as exc: + print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr) + return 1 + + from paddleocr import PaddleOCRVL + + print( + f"[PDF] {preflight['page_count']} pages, selected: " + f"{len(preflight['selected_pages'])}, output: {preflight['document_dir']}" + ) + print(f"[GPU] Device: {device} ({device_name})") + print("Loading PaddleOCR-VL-1.6...") + init_started = time.perf_counter() + pipeline = PaddleOCRVL(pipeline_version="v1.6", device=device) + paddle.device.cuda.synchronize(args.device_id) + init_seconds = time.perf_counter() - init_started + print(f"Model init: {init_seconds:.2f}s") + + predict_kwargs = { + key: value + for key, value in { + "max_new_tokens": args.max_new_tokens, + "min_pixels": args.min_pixels, + "max_pixels": args.max_pixels, + }.items() + if value is not None + } + metadata = { + "device": device, + "device_name": device_name, + "python_version": platform.python_version(), + "platform": platform.platform(), + "paddle_version": paddle.__version__, + "model_init_seconds": round(init_seconds, 3), + "pipeline_version": "v1.6", + } + + try: + summary = process_pdf( + pipeline=pipeline, + pdf_path=args.pdf, + output_root=args.output, + pages=args.pages, + dpi=args.dpi, + password=args.password, + resume=args.resume, + overwrite=args.overwrite, + keep_rendered=args.keep_rendered, + fail_fast=args.fail_fast, + run_metadata=metadata, + predict_kwargs=predict_kwargs, + synchronize=lambda: paddle.device.cuda.synchronize(args.device_id), + ) + except KeyboardInterrupt: + print("\n[INTERRUPTED] 已保存当前进度；使用 --resume 继续。", file=sys.stderr) + return 130 + except Exception as exc: + print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr) + return 1 + + print("\n[PDF OCR Summary]") + print(f"Status: {summary['status']}") + print(f"Completed: {summary['completed_pages']} / {summary['selected_pages']}") + print(f"Failed: {summary['failed_pages']}") + print(f"Output: {summary['document_dir']}") + return 0 if not summary["failed_pages"] else 3 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/gpu/pyproject.toml b/gpu/pyproject.toml index 733f289..97ae644 100644 --- a/gpu/pyproject.toml +++ b/gpu/pyproject.toml @@ -7,6 +7,7 @@ requires-python = ">=3.11,<3.13" dependencies = [ "paddleocr[doc-parser]==3.7.0", "paddlepaddle-gpu==3.2.1", + "pypdfium2>=5.11.0", "setuptools>=83.0.0", ] diff --git a/pdf_ocr.py b/pdf_ocr.py new file mode 100644 index 0000000..8fa9b83 --- /dev/null +++ b/pdf_ocr.py @@ -0,0 +1,127 @@ +"""CPU entry point for page-by-page PaddleOCR-VL PDF recognition.""" + +from __future__ import annotations + +import argparse +import os +import platform +import sys +import time +from pathlib import Path + +from pdf_ocr_core import preflight_pdf, process_pdf + +PROJECT_ROOT = Path(__file__).resolve().parent +DEFAULT_OUTPUT = PROJECT_ROOT / "outputs" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="PaddleOCR-VL-1.6 CPU PDF OCR（逐页、可恢复）", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("pdf", type=Path, help="输入 PDF 文件") + parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="输出根目录") + parser.add_argument("--pages", help="一页或多个页码范围，例如 1-5,8,10-") + parser.add_argument("--dpi", type=int, default=144, help="PDF 页面渲染 DPI") + parser.add_argument("--password", help="加密 PDF 密码") + parser.add_argument("--threads", type=int, default=None, help="Paddle CPU 线程数") + parser.add_argument("--resume", action="store_true", help="跳过已完成页，继续现有任务") + parser.add_argument("--overwrite", action="store_true", help="删除已有输出并重新处理") + parser.add_argument("--keep-rendered", action="store_true", help="保留逐页渲染 PNG") + parser.add_argument("--fail-fast", action="store_true", help="任一页失败后立即停止") + parser.add_argument("--max-new-tokens", type=int, default=None, help="限制每个文本块最大生成 token") + parser.add_argument("--min-pixels", type=int, default=None, help="VLM 最小输入像素参数") + parser.add_argument("--max-pixels", type=int, default=None, help="VLM 最大输入像素参数") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + total_cores = os.cpu_count() or 4 + safe_default_threads = max(1, total_cores - 2) + threads = args.threads or int(os.environ.get("PADDLE_THREADS", safe_default_threads)) + if threads < 1: + print("[ERROR] --threads 必须大于等于 1", file=sys.stderr) + return 2 + + try: + preflight = preflight_pdf( + pdf_path=args.pdf, + output_root=args.output, + pages=args.pages, + dpi=args.dpi, + password=args.password, + resume=args.resume, + overwrite=args.overwrite, + ) + except Exception as exc: + print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr) + return 1 + + print( + f"[PDF] {preflight['page_count']} pages, selected: " + f"{len(preflight['selected_pages'])}, output: {preflight['document_dir']}" + ) + + from paddle import core + from paddleocr import PaddleOCRVL + + core.set_num_threads(threads) + print(f"[CPU] Threads: {threads} / {total_cores} (reserved: {max(0, total_cores - threads)})") + print("Loading PaddleOCR-VL-1.6...") + init_started = time.perf_counter() + pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu") + init_seconds = time.perf_counter() - init_started + print(f"Model init: {init_seconds:.2f}s") + + predict_kwargs = { + key: value + for key, value in { + "max_new_tokens": args.max_new_tokens, + "min_pixels": args.min_pixels, + "max_pixels": args.max_pixels, + }.items() + if value is not None + } + metadata = { + "device": "cpu", + "cpu_threads": threads, + "python_version": platform.python_version(), + "platform": platform.platform(), + "model_init_seconds": round(init_seconds, 3), + "pipeline_version": "v1.6", + } + + try: + summary = process_pdf( + pipeline=pipeline, + pdf_path=args.pdf, + output_root=args.output, + pages=args.pages, + dpi=args.dpi, + password=args.password, + resume=args.resume, + overwrite=args.overwrite, + keep_rendered=args.keep_rendered, + fail_fast=args.fail_fast, + run_metadata=metadata, + predict_kwargs=predict_kwargs, + ) + except KeyboardInterrupt: + print("\n[INTERRUPTED] 已保存当前进度；使用 --resume 继续。", file=sys.stderr) + return 130 + except Exception as exc: + print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr) + return 1 + + print("\n[PDF OCR Summary]") + print(f"Status: {summary['status']}") + print(f"Completed: {summary['completed_pages']} / {summary['selected_pages']}") + print(f"Failed: {summary['failed_pages']}") + print(f"Output: {summary['document_dir']}") + return 0 if not summary["failed_pages"] else 3 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pdf_ocr_core.py b/pdf_ocr_core.py new file mode 100644 index 0000000..9ca39a4 --- /dev/null +++ b/pdf_ocr_core.py @@ -0,0 +1,534 @@ +"""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() diff --git a/pyproject.toml b/pyproject.toml index 87ea9b6..1b2673d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,8 @@ description = "Add your description here" readme = "README.md" requires-python = ">=3.13" dependencies = [ - "paddleocr[doc-parser]>=3.6.0", + "paddleocr[doc-parser]==3.7.0", "paddlepaddle==3.2.1", + "pypdfium2>=5.11.0", "setuptools>=83.0.0", ] diff --git a/uv.lock b/uv.lock index a7a4f72..c5d2c0d 100644 --- a/uv.lock +++ b/uv.lock @@ -1032,13 +1032,15 @@ source = { virtual = "." } dependencies = [ { name = "paddleocr", extra = ["doc-parser"] }, { name = "paddlepaddle" }, + { name = "pypdfium2" }, { name = "setuptools" }, ] [package.metadata] requires-dist = [ - { name = "paddleocr", extras = ["doc-parser"], specifier = ">=3.6.0" }, + { name = "paddleocr", extras = ["doc-parser"], specifier = "==3.7.0" }, { name = "paddlepaddle", specifier = "==3.2.1" }, + { name = "pypdfium2", specifier = ">=5.11.0" }, { name = "setuptools", specifier = ">=83.0.0" }, ]