From 2988521c410d0cccfcf80fcc16e824fdf6a90c5c Mon Sep 17 00:00:00 2001 From: kuuhaku Date: Thu, 16 Jul 2026 16:20:55 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=E8=B0=83=E6=95=B4=E4=BB=A3=E7=A0=81?= =?UTF-8?q?=E7=BB=93=E6=9E=84=EF=BC=8C=E5=B0=86cpu=E6=A8=A1=E5=BC=8F?= =?UTF-8?q?=E6=94=B6=E7=BA=B3=E8=BF=9B=E5=AD=90=E7=9B=AE=E5=BD=95=EF=BC=8C?= =?UTF-8?q?=E7=BB=9F=E4=B8=80=E5=85=A5=E5=8F=A3=EF=BC=8C=E6=A0=B9=E6=8D=AE?= =?UTF-8?q?=E6=96=87=E4=BB=B6=E5=90=8E=E7=BC=80=E5=92=8C=E7=9B=AE=E5=BD=95?= =?UTF-8?q?=E8=87=AA=E5=8A=A8=E5=88=A4=E6=96=AD=E4=BD=BF=E7=94=A8=E8=B7=AF?= =?UTF-8?q?=E7=94=B1=EF=BC=88pdf=E3=80=81=E5=9B=BE=E7=89=87=E5=92=8C?= =?UTF-8?q?=E6=89=B9=E9=87=8F=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 9 +- README.md | 726 +++++++++--------- batch_ocr.py | 329 -------- benchmarks/README.md | 28 +- benchmarks/cpu/.gitkeep | 0 .python-version => cpu/.python-version | 0 cpu/README.md | 16 + pyproject.toml => cpu/pyproject.toml | 11 +- cpu/runner.py | 14 + uv.lock => cpu/uv.lock | 55 +- .../documents}/化肥买卖合同　GF—2000—0102.pdf | Bin .../民用爆破器材买卖合同　GF—2001—0107.pdf | Bin {images => data/images}/名片01.jpg | Bin {images => data/images}/名片02.jpg | Bin {images => data/images}/手写01.png | Bin gpu/README.md | 148 +--- gpu/main.py | 291 ------- gpu/pdf_ocr.py | 214 ------ gpu/pyproject.toml | 4 +- gpu/runner.py | 14 + gpu/setup_env.py | 12 +- gpu/verify_env.py | 53 -- logs/{ => legacy}/名片01.log | 0 logs/{ => legacy}/名片02.log | 0 logs/{ => legacy}/手写01.log | 0 logs/{ => legacy}/批量识别.log | 0 main.py | 140 ---- ocr.py | 78 ++ ocr_app/__init__.py | 3 + ocr_app/cli.py | 124 +++ ocr_app/commands.py | 497 ++++++++++++ ocr_logging.py => ocr_app/logging_utils.py | 0 ocr_app/output.py | 139 ++++ ocr_app/pdf.py | 541 +++++++++++++ ocr_app/pdf_text.py | 127 +++ ocr_app/runtime.py | 170 ++++ pdf_ocr.py | 182 ----- pdf_ocr_core.py | 658 ---------------- tests/conftest.py | 6 + tests/test_launcher.py | 13 + tests/test_output_routing.py | 126 +++ tests/test_page_spec.py | 13 + tests/test_pdf_hybrid.py | 102 +++ tests/test_pdf_text.py | 39 + 44 files changed, 2522 insertions(+), 2360 deletions(-) delete mode 100644 batch_ocr.py create mode 100644 benchmarks/cpu/.gitkeep rename .python-version => cpu/.python-version (100%) create mode 100644 cpu/README.md rename pyproject.toml => cpu/pyproject.toml (53%) create mode 100644 cpu/runner.py rename uv.lock => cpu/uv.lock (99%) rename {documents => data/documents}/化肥买卖合同　GF—2000—0102.pdf (100%) rename {documents => data/documents}/民用爆破器材买卖合同　GF—2001—0107.pdf (100%) rename {images => data/images}/名片01.jpg (100%) rename {images => data/images}/名片02.jpg (100%) rename {images => data/images}/手写01.png (100%) delete mode 100644 gpu/main.py delete mode 100644 gpu/pdf_ocr.py create mode 100644 gpu/runner.py delete mode 100644 gpu/verify_env.py rename logs/{ => legacy}/名片01.log (100%) rename logs/{ => legacy}/名片02.log (100%) rename logs/{ => legacy}/手写01.log (100%) rename logs/{ => legacy}/批量识别.log (100%) delete mode 100644 main.py create mode 100644 ocr.py create mode 100644 ocr_app/__init__.py create mode 100644 ocr_app/cli.py create mode 100644 ocr_app/commands.py rename ocr_logging.py => ocr_app/logging_utils.py (100%) create mode 100644 ocr_app/output.py create mode 100644 ocr_app/pdf.py create mode 100644 ocr_app/pdf_text.py create mode 100644 ocr_app/runtime.py delete mode 100644 pdf_ocr.py delete mode 100644 pdf_ocr_core.py create mode 100644 tests/conftest.py create mode 100644 tests/test_launcher.py create mode 100644 tests/test_output_routing.py create mode 100644 tests/test_page_spec.py create mode 100644 tests/test_pdf_hybrid.py create mode 100644 tests/test_pdf_text.py diff --git a/.gitignore b/.gitignore index 1b28b4f..35cac98 100644 --- a/.gitignore +++ b/.gitignore @@ -8,15 +8,22 @@ wheels/ # Virtual environments .venv +cpu/.venv/ +gpu/.venv/ +gpu/.gpu-ready # Generated benchmark results +benchmarks/cpu/*.json benchmarks/gpu/*.json +!benchmarks/cpu/.gitkeep !benchmarks/gpu/.gitkeep # OCR outputs outputs/ # Generated structured logs (legacy logs directly under logs/ remain tracked) -logs/single/ +logs/input/ +logs/verify/ +logs/image/ logs/batch/ logs/pdf/ diff --git a/README.md b/README.md index 6e513d9..1a8491b 100644 --- a/README.md +++ b/README.md @@ -1,428 +1,420 @@ -# ocr-VL1.6 +# PaddleOCR-VL-1.6 本地 OCR -本地部署 [PaddlePaddle/PaddleOCR-VL-1.6](https://github.com/PaddlePaddle/PaddleOCR) 的实验项目，包含已实测的 CPU 版本和独立隔离、待 NVIDIA GPU 验证的 GPU 版本。 +本项目使用 PaddleOCR-VL-1.6 实现统一的图片 OCR、批量图片 OCR 和 PDF 识别，CPU/GPU 环境完全隔离，并通过根目录唯一入口 `ocr.py` 调用。 -> 在线 Demo: [HuggingFace Space](https://huggingface.co/spaces/PaddlePaddle/PaddleOCR-VL-1.6_Online_Demo) · 模型权重: [HuggingFace](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.6) +PDF 默认使用 **文本提取 + OCR 混合模式**：优先提取 PDF 原始文本层，仅当页面没有有效文本层或文本质量不足时才加载 PaddleOCR-VL 并 OCR。 -## 项目结构 +> 当前开发机器只有集成显卡。CPU 功能已验证；GPU 代码已实现，但必须在 NVIDIA CUDA GPU 机器上安装和测试。 -``` +## 目录结构 + +```text ocr-VL1.6/ -├── main.py # CPU 单图 OCR + Benchmark -├── batch_ocr.py # CPU 批量图片 OCR（系统友好的多进程版本） -├── pdf_ocr.py # CPU PDF OCR（逐页、可恢复） -├── pdf_ocr_core.py # CPU/GPU 共用的 PDF 渲染、恢复和导出逻辑 -├── ocr_logging.py # CPU/GPU 共用的 UTF-8 结构化日志工具 -├── 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 独立依赖 -│ ├── .python-version # GPU 使用 Python 3.11 -│ └── README.md # GPU 安装与运行说明 -├── benchmarks/ -│ └── gpu/ # GPU Benchmark JSON 输出目录 -├── images/ -│ └── 手写01.png # 测试图片：手写中文（1758×646） -└── README.md +├── ocr.py # 唯一用户入口，自动选择 CPU/GPU 子环境 +├── ocr_app/ # CPU/GPU 共享业务代码 +│ ├── cli.py # image/batch/pdf/verify 命令 +│ ├── commands.py # 命令实现 +│ ├── runtime.py # 设备验证、模型延迟加载 +│ ├── pdf.py # 混合 PDF、断点续传、导出 +│ ├── pdf_text.py # 文本层提取与质量评估 +│ └── logging_utils.py # UTF-8 结构化日志 +├── cpu/ +│ ├── pyproject.toml # CPU 独立依赖 +│ ├── uv.lock # CPU 独立锁文件 +│ ├── .python-version # Python 3.13 +│ └── runner.py # 统一入口的 CPU 执行器 +├── gpu/ +│ ├── pyproject.toml # GPU 独立依赖 +│ ├── .python-version # Python 3.11 +│ ├── setup_env.py # CUDA Wheel 安装脚本 +│ └── runner.py # 统一入口的 GPU 执行器 +├── data/ +│ ├── images/ # 测试图片 +│ └── documents/ # 测试 PDF +├── outputs/ # PDF 输出（git ignored） +├── benchmarks/ # 单图 Benchmark JSON +├── logs/ # UTF-8 运行日志 +└── tests/ # 共享逻辑测试 ``` -## 技术栈 +根目录不再保存 Paddle 虚拟环境或 Python 项目依赖。CPU 与 GPU 分别使用 `cpu/.venv` 和 `gpu/.venv`。 -| 组件 | 版本 | 说明 | -| ---------------- | ----- | ---------------------------------------- | -| Python（CPU） | 3.13 | 根目录独立环境 | -| Python（GPU） | 3.11 | `gpu/` 独立环境，提升 GPU Wheel 兼容性 | -| PaddlePaddle | 3.2.1 | CPU 使用 `paddlepaddle`，GPU 使用 `paddlepaddle-gpu` | -| PaddleOCR | 3.7.0 | 带 `doc-parser` extra | -| PaddleOCR-VL-1.6 | 0.9B | 主 OCR 视觉语言模型（~1.8GB） | -| PP-DocLayoutV3 | - | 版面检测模型（~126MB） | +## 安装 -模型缓存目录：`~/.paddlex/official_models/` - -## 快速开始 - -### 前提条件 - -- Python ≥ 3.13 -- [uv](https://github.com/astral-sh/uv) 包管理器 - -### 安装 +### CPU ```bash -uv sync +uv sync --project cpu ``` -### 运行 +### GPU + +根据目标机器 CUDA/驱动和 PaddlePaddle 官方兼容表选择 Wheel： ```bash -# 单张图片 OCR（自动使用全部 CPU 核心） -uv run python main.py - -# 批量 OCR（多进程并行，安全默认值） -uv run python batch_ocr.py images/ -``` - -所有 OCR 入口默认同时输出控制台日志和 UTF-8 日志文件，详见“运行日志”章节。 - -首次运行会自动从 ModelScope 下载模型文件（约 2GB），后续使用缓存。 - -### GPU 子项目 - -> **状态：已实现、未实测。** 当前开发机器只有集成显卡，不能运行 NVIDIA CUDA。GPU 代码已通过语法、CLI 和无 CUDA 安全退出检查，但安装兼容性、显存占用和性能数据必须在目标 NVIDIA GPU 机器上验证。 - -CPU 与 GPU 使用不同虚拟环境，禁止在根目录 CPU `.venv` 中安装 `paddlepaddle-gpu`。 - -```bash -# 查看 GPU 安装命令，不实际安装 python gpu/setup_env.py --cuda cu118 --dry-run - -# 在目标 NVIDIA GPU 机器创建 gpu/.venv；根据官方兼容表选择 cu118 或 cu126 python gpu/setup_env.py --cuda cu118 - -# 检查 CUDA 构建、GPU 设备和矩阵乘法 -uv run --project gpu python gpu/verify_env.py - -# 运行 GPU 单图 Benchmark -uv run --project gpu python gpu/main.py --warmup 1 --rounds 3 ``` -GPU Benchmark JSON 写入 `benchmarks/gpu/`。详细说明见 [`gpu/README.md`](gpu/README.md)。 - -## 运行日志 - -所有主要入口均使用统一日志格式： - -```text -2026-07-16 14:28:02 | INFO | pid=27644 | PAGE_OCR_COMPLETED page=1 seconds=36.345 -``` - -默认日志目录： - -```text -logs/ -├── single/ # main.py / gpu/main.py -├── batch/ # batch_ocr.py -└── pdf/ # pdf_ocr.py / gpu/pdf_ocr.py -``` - -默认文件名包含输入名、设备和时间戳，例如： - -```text -logs/pdf/sample-cpu-20260716-142802.log -logs/single/手写01-gpu0-20260716-142802.log -``` - -可用参数： +也支持： ```bash -# 指定日志文件 -uv run python main.py images/手写01.png --log-file logs/custom.log -uv run python pdf_ocr.py documents/sample.pdf --log-file logs/pdf-sample.log -uv run python batch_ocr.py images/ --log-file logs/batch-images.log - -# 输出详细异常堆栈和调试日志 -uv run python pdf_ocr.py documents/sample.pdf --verbose +python gpu/setup_env.py --cuda cu126 ``` -日志文件使用 UTF-8 编码。即使 Windows 控制台因 GBK 显示乱码，日志文件中的中文仍可正常查看。 +安装成功后脚本生成 `gpu/.gpu-ready`。统一入口只会调用已安装的 `gpu/.venv`，不会从默认 PyPI 误装 GPU 包，也不会回退到 CPU。 -### 单图日志统计 +## 简化统一入口 -`main.py` 与 `gpu/main.py` 记录： - -- 程序启动与输入图片大小 -- Paddle/PaddleOCR 导入耗时 -- CPU 线程数或 GPU/CUDA 初始化耗时 -- 模型初始化耗时 -- 每轮预热耗时 -- 每轮正式推理耗时 -- min/max/mean/median/stdev -- 图片尺寸、版面框数量、文本块数量 -- OCR 文本块内容（可用 `--no-result` 关闭） -- 从程序启动到结果输出的总用时 -- GPU 入口额外记录显存统计和 Benchmark JSON 路径 - -### 批量图片日志统计 - -`batch_ocr.py` 记录： - -- 图片扫描耗时和图片数量 -- Worker 数、每 Worker 线程数和预估内存 -- 每个 Worker 的 PID、错峰等待、框架导入、模型初始化和启动总耗时 -- 每张图片的 Worker PID、推理耗时、尺寸、版面框和文本块数量 -- 任务进度、成功数和失败数 -- Pool 总耗时、串行耗时估计、平均每图耗时和并行加速比 -- 从程序启动到全部结果汇总的总用时 - -### PDF 日志统计 - -`pdf_ocr.py` 与 `gpu/pdf_ocr.py` 记录： - -- PDF 预检、打开和 manifest 创建耗时 -- 模型初始化耗时 -- 每页渲染耗时 -- 每页 OCR 推理耗时 -- 每页 Markdown/JSON 导出耗时 -- manifest 与合并文件保存耗时 -- 每页总耗时、累计耗时和预计剩余时间（ETA） -- 每页图片尺寸、版面框数量和文本块数量 -- 完成页、失败页和断点续传前已完成页数 -- 各阶段累计值、平均每页耗时和任务总用时 -- 从程序启动（含模型加载）到退出的程序总用时 - -PDF 的 `manifest.json` 同时包含 `summary.timing`： - -```json -{ - "pdf_open_seconds": 0.01, - "manifest_prepare_seconds": 0.03, - "render_total_seconds": 1.2, - "ocr_total_seconds": 324.5, - "export_total_seconds": 0.8, - "state_save_total_seconds": 0.2, - "page_total_seconds": 326.7, - "average_ocr_seconds": 162.25, - "average_page_seconds": 163.35, - "finalize_seconds": 0.1, - "task_total_seconds": 327.1 -} -``` - -`task_total_seconds` 是 PDF 核心任务总时间，不含入口模型初始化；完整程序总时间记录在日志的 `PROGRAM_COMPLETED` 事件中。 - -## 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 +python ocr.py <文件或目录> --device cpu|gpu ``` -### GPU 使用 +程序自动判断输入类型： -GPU 入口与 CPU 入口使用相同的 PDF 核心逻辑，但必须在 `gpu/.venv` 和 NVIDIA CUDA GPU 上运行： +| 输入 | 自动路由 | +|------|----------| +| `.png/.jpg/.jpeg/.bmp/.tif/.tiff/.webp` | 图片 OCR + Benchmark | +| `.pdf` | PDF 混合文本提取/OCR | +| 目录 | 发现其中的图片和 PDF，逐个调用同一个单文件路由器 | + +不支持的文件后缀会给出明确错误；目录模式会自动忽略不支持的文件。 + +### 验证环境 ```bash -uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf \ - --device-id 0 \ - --pages "1-10" \ - --dpi 144 +python ocr.py verify --device cpu +python ocr.py verify --device gpu ``` -当前机器无 NVIDIA 独立显卡，因此 GPU PDF 入口仅完成静态检查，尚未实机验证。 +### 单张图片 -### 页码语法 +```bash +python ocr.py data/images/手写01.png --device cpu +``` -| 参数 | 含义 | -|------|------| -| `1` | 仅第 1 页 | -| `1-5` | 第 1～5 页 | -| `10-` | 第 10 页到最后一页 | -| `1-5,8,10-` | 多个页码范围组合 | +多轮 Benchmark： -用户页码从 1 开始；内部 manifest 使用同样的一基页码记录。 +```bash +python ocr.py data/images/手写01.png \ + --device cpu \ + --warmup 1 \ + --rounds 3 +``` -### 输出结构 +图片识别结果和 Benchmark 默认一起写入： + +```text +outputs/images/<图片名_扩展名>/ +``` + +### 单个 PDF + +```bash +python ocr.py data/documents/sample.pdf --device cpu +``` + +PDF 默认使用混合模式。强制模式： + +```bash +python ocr.py sample.pdf --pdf-mode text --device cpu +python ocr.py sample.pdf --pdf-mode ocr --device cpu +``` + +`--mode` 仍作为 `--pdf-mode` 的简写别名保留。 + +### 批量目录 + +```bash +# 扫描当前目录层级中的图片和 PDF +python ocr.py data/ --device cpu + +# 递归扫描所有子目录 +python ocr.py data/ --recursive --device cpu +``` + +目录模式的底层就是重复调用同一个单文件路由器，因此： + +- 图片使用与单图片相同的 Benchmark 和日志逻辑 +- PDF 使用与单 PDF 相同的混合路由、断点和导出逻辑 +- 所有文件共享同一个延迟加载模型实例 +- 电子 PDF 不会触发模型加载；首张图片或首个 OCR 页面才加载模型 +- 保持串行处理，避免此前多进程造成卡顿、无响应和黑屏 + +递归目录的图片和 PDF 输出都会保留相对目录结构，避免不同子目录下同名文件相互覆盖。 + +## 统一输出目录 + +图片和 PDF 现在都会写入 `--output` 指定目录，默认是 `outputs/`： ```text outputs/ -└── sample/ - ├── manifest.json # 任务配置、页状态、耗时和错误 - ├── document.md # 合并后的 Markdown - ├── document.json # 合并后的 JSON - ├── pages/ - │ ├── page-0001.md - │ ├── page-0001.json - │ └── ... - ├── assets/ # 表格、图片等 Markdown 资源 - └── rendered/ # 仅使用 --keep-rendered 时保留 +├── images/ +│ └── <相对目录>/<图片名_扩展名>/ +│ ├── result.md # Markdown 结果 +│ ├── result.txt # 纯文本结果 +│ ├── result.json # PaddleOCR 结构化结果 +│ ├── benchmark.json # 模型/推理/导出耗时 +│ └── assets/ # 识别结果中的图片资源（如有） +├── pdfs/ +│ └── <相对目录>// +│ ├── manifest.json +│ ├── document.md +│ ├── document.json +│ └── pages/ +└── batches/ + └── <目录名>-<时间戳>.json # 批量任务汇总 ``` -默认不保留中间渲染 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 分两阶段： - -``` -输入图片 → [PP-DocLayoutV3 版面检测] → [PaddleOCR-VL-1.6-0.9B 文字识别] → 结构化输出 -``` - -1. **版面检测** — PP-DocLayoutV3 检测页面中的文本块区域（坐标 + 标签 + 置信度） -2. **OCR 识别** — PaddleOCR-VL-1.6-0.9B（GQA 架构视觉语言模型）逐块识别文字 -3. **结果输出** — 返回 `PaddleOCRVLResult`，包含布局信息和识别文本 - -### 输出结构 - -| 字段 | 类型 | 说明 | -| ---------------------- | ---------------------- | ------------------------------------------------------------ | -| `layout_det_res.boxes` | list[dict] | 版面文本区域（cls_id, label, score, coordinate, polygon_points） | -| `parsing_res_list` | list[PaddleOCRVLBlock] | 识别文本块（$.label, $.bbox, $.content, $.polygon_points） | -| `model_settings` | dict | 推理配置开关（版面检测/图表/印章等） | -| `width` / `height` | int | 图片尺寸 | - -## 性能优化迭代 - -测试机器：Windows 11, CPU 20 核（逻辑线程）, RAM 32GB, PaddlePaddle 3.2.1 CPU - -### 迭代 0：初始状态（无任何优化） - -直接调用 `pipeline.predict()`，未设置任何线程参数。 - -| 阶段 | 耗时 | -| ----------------------- | --------------- | -| 模型初始化（加载权重） | ~60s | -| 首次推理（含 JIT 编译） | ~285s | -| 后续推理 | ~238s（~4 min） | - -### 迭代 1：算子级并行 — `core.set_num_threads()` - -**探索过程：** - -| 尝试 | 方法 | 结果 | -| ---- | -------------------------------------------------- | ---------------------------------------------------- | -| ❌ | `paddle.set_num_threads()` | Paddle 3.x 已移除该 API | -| ❌ | 环境变量 `OMP_NUM_THREADS` / `MKL_NUM_THREADS` | Paddle 3.x 内部使用 oneDNN，不读取这些变量 | -| ✅ | `from paddle import core; core.set_num_threads(N)` | **有效！** oneDNN 底层算子（matmul 等）受该 API 控制 | - -**矩阵乘法微基准测试（4000×4000）：** - -| 线程数 | 耗时 (matmul) | 加速比 | -| ------ | ------------- | -------- | -| 1 | 0.952s | 1.0x | -| 4 | 0.419s | 2.3x | -| 8 | 0.323s | 2.9x | -| 16 | 0.240s | 4.0x | -| **20** | **0.223s** | **4.3x** | - -**应用到 OCR 后的实际效果：** - -设置 `core.set_num_threads(20)` 后重新评测： - -| 阶段 | 优化前 | 优化后 | 提速 | -| ---------- | ------ | --------------------- | -------- | -| 模型初始化 | ~60s | ~40s | 1.5x | -| 推理 | ~238s | **~162s（~2.7 min）** | **1.5x** | - -**为什么不是 4.3x？** 矩阵乘法只是 OCR pipeline 的一部分。自回归解码（逐 token 生成）天然串行、I/O 等待、版面检测中的非矩阵运算等不受线程数影响。 - ---- - -### 迭代 2：批量多进程并行 — `batch_ocr.py` - -**思路：** 多张图片时，用 `multiprocessing.Pool` 启动多个独立进程，每个进程加载一份 pipeline 实例，同时处理不同图片。 - -**遇到的问题 & 修复（迭代 2.1）：** - -| 问题 | 原因 | 修复 | -|------|------|------| -| 系统卡顿/黑屏/无响应 | `Pool.starmap` 同时启动 N 个进程，同步加载 N×2GB 模型，CPU + 内存瞬间打满 | ① 进程错峰启动（随机延迟 0~15s）② `psutil` 降低进程优先级 ③ 预留 1 核给 OS ④ `imap_unordered` 替代 `starmap` | - -**策略：** -- 每个子进程独立调用 `core.set_num_threads((总核心-1) / 进程数)`，预留核心给 OS -- 例如 2 进程 × 9 线程 = 18 核，留 2 核给系统 -- `--stagger` 控制错峰窗口，默认 15s +例如： ```bash -# 2 进程并行（安全默认值） -uv run python batch_ocr.py images/ - -# 4 进程并行（需 32GB+ RAM） -uv run python batch_ocr.py images/ --workers 4 - -# 自定义错峰窗口（值越大内存峰值越低，但总耗时增加） -uv run python batch_ocr.py images/ --workers 4 --stagger 30 +python ocr.py data/images/名片01.jpg --device cpu ``` -| 配置 | 适用场景 | 理论加速比 | 内存开销 | 实际限制 | -| -------------------- | -------- | --------------- | ---------- | -------------- | -| `set_num_threads(N)` | 单张图片 | ~1.5x | 无额外开销 | 自回归解码瓶颈 | -| `batch_ocr.py` | 批量多图 | ~Nx（N=进程数） | N × 2GB | 内存/内存带宽，需错峰避免打满系统 | +会生成： -> ⚠️ 每个进程独立加载模型（~2GB），32GB RAM 建议从 `--workers 2` 开始测试。默认值是相对保守配置，但是否稳定仍取决于可用内存、散热、后台应用和图片复杂度。 +```text +outputs/images/名片01_jpg/result.md +outputs/images/名片01_jpg/result.txt +outputs/images/名片01_jpg/result.json +outputs/images/名片01_jpg/benchmark.json +``` ---- +图片目录名包含扩展名，避免 `same.png` 与 `same.jpg` 相互覆盖。 -### 迭代 3：独立 GPU 子项目（待实机验证） +重构前生成的 PDF 结果可能仍直接位于 `outputs//`。这些旧结果不会自动删除；新任务统一写入 `outputs/pdfs//`，确认不再需要后可手动迁移或清理。 -为避免 `paddlepaddle` 和 `paddlepaddle-gpu` 相互覆盖，在同一仓库新增 `gpu/` 子项目，使用独立 Python、虚拟环境、依赖配置和锁文件。 +目录任务再次运行时： -已完成： +- 已存在的 PDF manifest 自动断点续传 +- 新加入的 PDF 自动创建新任务 +- 图片重新识别，并使用原子写入覆盖对应输出文件 +- `--overwrite` 会强制 PDF 重新处理 +- 每次目录任务都会生成新的 `outputs/batches/*.json` 汇总 -- `gpu/setup_env.py`：根据 `cu118` / `cu126` Wheel 索引创建环境 -- `gpu/verify_env.py`：检查 CUDA 构建、设备数量并执行 GPU 矩阵乘法 -- `gpu/main.py`：显式指定 `device="gpu:N"`，支持预热、多轮计时和 JSON 输出 -- 无 CUDA 时立即退出，不静默回退到 CPU -- CPU 环境下已通过 Python 语法、CLI 和安全退出检查 +旧命令前缀 `image`、`pdf`、`batch` 暂时兼容，例如 `python ocr.py image a.png`，但推荐直接传入路径。 -尚未验证： +## PDF 混合模式 -- NVIDIA 驱动、CUDA Wheel 与目标 GPU 的兼容性 -- PaddleOCR-VL-1.6 GPU 模型初始化是否正常 -- GPU 显存峰值和真实推理速度 -- FP16/BF16、TensorRT 或批量推理收益 +### 默认：hybrid ---- +```bash +python ocr.py data/documents/sample.pdf --device cpu +``` -### 优化总结 +每页流程： -| 迭代 | 状态 | 结果 | -|------|------|------| -| CPU 初始版本 | 已实测 | 后续单图约 238s | -| CPU `set_num_threads(20)` | 已实测 | 单图约 162s，约 1.5x 加速 | -| CPU 多进程批量 | 已实现，稳定性依机器而定 | 理论提升批量吞吐；当前没有足够的可靠实测数据支持固定加速比 | -| 独立 GPU 子项目 | 已实现，未实机验证 | 等待 NVIDIA CUDA GPU 测试 | +```text +读取 PDF 文本层 + ↓ +文本质量评估 + ┌────┴────┐ + │ │ +有效文本 无效/不足 + │ │ +直接保存 渲染页面 → PaddleOCR-VL + └────┬────┘ + ↓ +逐页 Markdown/JSON + 合并文档 +``` -> CPU 单图当前实测约 2.7 分钟。批量多进程主要提高总吞吐，不会缩短某一张图片自身的推理延迟。GPU 性能在实机验证前不作预测。 +如果整份 PDF 都具有有效文本层，模型完全不会加载。实测项目中的电子合同 PDF，单页混合处理约 `0.06s`，且 `model_used=false`。 -## 已知局限 +### 强制文本模式 -| 问题 | 影响 | 说明 | -| ------------------ | ----------------------- | ------------------------------------- | -| CPU 推理极慢 | 单图 ~2.7 min（优化后） | 0.9B VL 模型不适合 CPU 实时场景 | -| 自回归解码串行 | 无法更细粒度并行 | 生成阶段逐 token 依赖，多线程收益有限 | -| 内存占用大 | 每进程需 ~2GB | 限制了 `batch_ocr.py` 并行度 | -| Windows 控制台乱码 | 中文输出显示为乱码 | GBK 编码问题，文件写入/pipe 正常 | -| GPU 未实机验证 | 暂无 GPU 性能结论 | 当前机器只有集成显卡，需 NVIDIA CUDA GPU 验证 | -| ccache 警告 | 无实际影响 | 仅影响首次编译加速，可忽略 | +```bash +python ocr.py data/documents/sample.pdf \ + --device cpu \ + --pdf-mode text +``` + +所有页面只提取 PDF 文本层，不执行 OCR。扫描页可能得到空文本。 + +### 强制 OCR 模式 + +```bash +python ocr.py data/documents/sample.pdf \ + --device cpu \ + --pdf-mode ocr +``` + +所有页面都渲染后交给 PaddleOCR-VL，适合模型一致性 Benchmark。 + +## 文本层质量判定 + +混合模式默认要求： + +| 参数 | 默认值 | 含义 | +|------|-------:|------| +| `--text-min-chars` | 50 | 非空白字符最小数量 | +| `--text-min-printable-ratio` | 0.85 | 可打印字符最低比例 | +| `--text-min-content-ratio` | 0.60 | 字母、数字、CJK 字符最低比例 | +| `--text-max-replacement-ratio` | 0.02 | Unicode 替换字符最高比例 | +| `--text-min-density` | 25 | 页面文本密度最低值 | + +例如某些扫描 PDF 只有页码或隐藏乱码层，混合模式会因为 `too_few_characters`、`low_content_ratio` 或 `high_replacement_ratio` 自动回退 OCR。 + +每页的判定结果会写入日志和 manifest： + +```json +{ + "source_type": "ocr", + "routing_reason": "too_few_characters", + "text_layer": { + "usable": false, + "non_whitespace_chars": 8, + "printable_ratio": 1.0, + "content_ratio": 0.75 + } +} +``` + +## PDF 页码与恢复 + +```bash +# 第 1～5 页、第 8 页、第 10 页到末页 +python ocr.py sample.pdf --pages "1-5,8,10-" --device cpu + +# 中断后继续 +python ocr.py sample.pdf --resume --device cpu + +# 删除旧输出并重做 +python ocr.py sample.pdf --overwrite --device cpu +``` + +`--resume` 会校验： + +- PDF SHA-256 +- PDF 处理模式 +- DPI +- 文本层判定阈值 +- manifest 版本 + +旧纯 OCR manifest（版本 1）不能直接用于新混合模式，请使用 `--overwrite` 重新生成。 + +## PDF 输出 + +```text +outputs/pdfs// +├── manifest.json +├── document.md +├── document.json +├── pages/ +│ ├── page-0001.md +│ ├── page-0001.json +│ └── ... +├── assets/ +└── rendered/ # 仅 --keep-rendered 时存在 +``` + +合并 Markdown 标题会标明页面来源： + +```markdown +## Page 1 (text) +## Page 2 (ocr) +``` + +Manifest 汇总包括： + +```json +{ + "text_pages": 8, + "ocr_pages": 2, + "model_used": true, + "model_initialized_during_task": true, + "timing": { + "text_extract_total_seconds": 0.15, + "render_total_seconds": 0.8, + "ocr_total_seconds": 324.5, + "model_init_seconds": 40.0, + "task_total_seconds": 366.0 + } +} +``` + +## PDF 常用参数 + +```bash +python ocr.py sample.pdf \ + --device cpu \ + --pdf-mode hybrid \ + --pages "1-10" \ + --dpi 144 \ + --threads 18 \ + --keep-rendered \ + --log-file logs/pdf-sample.log +``` + +DPI 建议： + +| 文档 | DPI | +|------|----:| +| 普通打印文字 | 120～144 | +| 小字号 | 150～200 | +| 手写/低质量扫描件 | 200～250 | + +DPI 只影响需要 OCR 的页面，不影响直接提取文本层的页面。 + +## 日志 + +统一日志格式： + +```text +2026-07-16 15:07:45 | INFO | pid=33552 | PAGE_ROUTED page=1 source=text reason=usable_text_layer +``` + +默认目录： + +```text +logs/input/ +logs/verify/ +logs/legacy/ +``` + +主要事件： + +- `RUNTIME_PREPARED` +- `MODEL_INITIALIZED` +- `FILE_ROUTED` +- `PAGE_ROUTED` +- `PAGE_FINISHED` +- `TASK_COMPLETED` +- `IMAGE_COMPLETED` +- `DIRECTORY_SUMMARY` +- `VERIFY_COMPLETED` + +日志使用 UTF-8。Windows 控制台即使显示乱码，日志文件中的中文仍正常。 + +## 测试 + +```bash +uv run --project cpu pytest -q +``` + +当前测试覆盖： + +- 页码范围解析 +- 文本标准化 +- 有效文本层判定 +- 空文本/短文本自动回退条件 +- 根入口设备解析 + +当前结果： + +```text +17 passed +``` + +## 已验证状态 + +- CPU `verify`：通过 +- 单图统一入口：假模型端到端通过，原真实模型能力此前已验证 +- 批量统一入口：假模型端到端通过 +- PDF `hybrid` 电子文本页：真实 PDF 验证通过，未加载模型 +- PDF 扫描页自动回退 OCR：假模型端到端通过 +- PDF `text` / `ocr` 强制模式：通过 +- GPU 环境路由：无环境时安全提示，不回退 CPU +- GPU 实机推理：尚未验证 diff --git a/batch_ocr.py b/batch_ocr.py deleted file mode 100644 index 4f337f3..0000000 --- a/batch_ocr.py +++ /dev/null @@ -1,329 +0,0 @@ -"""System-friendly multiprocessing batch OCR with structured timing logs.""" - -from __future__ import annotations - -import argparse -import logging -import os -import random -import sys -import time -from logging.handlers import QueueHandler, QueueListener -from multiprocessing import Manager, Pool, cpu_count -from pathlib import Path - -from ocr_logging import default_log_path, setup_run_logger - -PROJECT_ROOT = Path(__file__).resolve().parent -_WORKER_LOG_QUEUE = None -_WORKER_INIT_METRICS: dict = {} - - -def _worker_logger() -> logging.Logger: - logger = logging.getLogger(f"ocr.batch.worker.{os.getpid()}") - if logger.handlers: - return logger - logger.setLevel(logging.INFO) - logger.propagate = False - if _WORKER_LOG_QUEUE is not None: - logger.addHandler(QueueHandler(_WORKER_LOG_QUEUE)) - return logger - - -def _init_worker(threads: int, stagger_max: float, log_queue) -> None: - """Stagger startup, lower process priority, and load one model per worker.""" - global _pipeline, _WORKER_LOG_QUEUE, _WORKER_INIT_METRICS - _WORKER_LOG_QUEUE = log_queue - logger = _worker_logger() - worker_started = time.perf_counter() - delay = random.uniform(0, stagger_max) - logger.info("WORKER_START threads=%d stagger_delay_seconds=%.3f", threads, delay) - time.sleep(delay) - - import_started = time.perf_counter() - from paddle import core - core.set_num_threads(threads) - import_seconds = time.perf_counter() - import_started - - try: - import psutil - - process = psutil.Process() - if sys.platform == "win32": - process.nice(psutil.BELOW_NORMAL_PRIORITY_CLASS) - else: - process.nice(10) - priority_status = "lowered" - except Exception as exc: - priority_status = f"unchanged:{type(exc).__name__}" - - model_started = time.perf_counter() - from paddleocr import PaddleOCRVL - - _pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu") - model_seconds = time.perf_counter() - model_started - startup_total = time.perf_counter() - worker_started - _WORKER_INIT_METRICS = { - "pid": os.getpid(), - "threads": threads, - "stagger_delay_seconds": round(delay, 3), - "import_seconds": round(import_seconds, 3), - "model_init_seconds": round(model_seconds, 3), - "startup_total_seconds": round(startup_total, 3), - "priority": priority_status, - } - logger.info( - "WORKER_READY threads=%d import_seconds=%.3f model_init_seconds=%.3f startup_total_seconds=%.3f priority=%s", - threads, - import_seconds, - model_seconds, - startup_total, - priority_status, - ) - - -def _ocr_task(image_path: str) -> dict: - global _pipeline, _WORKER_INIT_METRICS - logger = _worker_logger() - started = time.perf_counter() - logger.info("IMAGE_START path=%s", image_path) - try: - result = _pipeline.predict(image_path) - elapsed = time.perf_counter() - started - first = result[0] - blocks = [ - {"label": block.label, "bbox": block.bbox, "content": block.content} - for block in first["parsing_res_list"] - if block.content.strip() - ] - response = { - "path": image_path, - "status": "completed", - "elapsed": round(elapsed, 3), - "width": first.get("width"), - "height": first.get("height"), - "layout_boxes": len(first["layout_det_res"]["boxes"]), - "parsed_blocks": len(first["parsing_res_list"]), - "blocks": blocks, - "worker_pid": os.getpid(), - "worker_init": _WORKER_INIT_METRICS, - } - logger.info( - "IMAGE_COMPLETED path=%s seconds=%.3f width=%s height=%s layout_boxes=%d parsed_blocks=%d non_empty_blocks=%d", - image_path, - elapsed, - response["width"], - response["height"], - response["layout_boxes"], - response["parsed_blocks"], - len(blocks), - ) - return response - except Exception as exc: - elapsed = time.perf_counter() - started - logger.exception("IMAGE_FAILED path=%s seconds=%.3f error=%s", image_path, elapsed, exc) - return { - "path": image_path, - "status": "failed", - "elapsed": round(elapsed, 3), - "error": f"{type(exc).__name__}: {exc}", - "blocks": [], - "worker_pid": os.getpid(), - "worker_init": _WORKER_INIT_METRICS, - } - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="批量 OCR — 多进程并行（系统友好版）", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument("dir", type=Path, help="图片目录") - parser.add_argument("--workers", type=int, default=2, help="并行进程数") - parser.add_argument("--threads", type=int, default=None, help="每进程线程数") - parser.add_argument("--stagger", type=float, default=15.0, help="Worker 启动错峰窗口秒数") - parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") - parser.add_argument("--verbose", action="store_true", help="输出详细日志") - parser.add_argument("--no-result", action="store_true", help="不记录 OCR 文本块") - return parser.parse_args() - - -def main() -> int: - program_started = time.perf_counter() - args = parse_args() - image_dir = args.dir.expanduser().resolve() - log_file = args.log_file or default_log_path(PROJECT_ROOT, "batch", image_dir.name, device="cpu") - logger = setup_run_logger("ocr.batch.main", log_file, verbose=args.verbose) - - if not image_dir.is_dir(): - logger.error("INPUT_DIRECTORY_NOT_FOUND path=%s", image_dir) - return 1 - if args.workers < 1 or args.stagger < 0: - logger.error("INVALID_ARGUMENT workers=%d stagger=%.3f", args.workers, args.stagger) - return 2 - - scan_started = time.perf_counter() - extensions = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tiff", "*.tif", "*.webp") - images = sorted(path for extension in extensions for path in image_dir.glob(extension)) - scan_seconds = time.perf_counter() - scan_started - if not images: - logger.error("NO_IMAGES_FOUND path=%s scan_seconds=%.3f", image_dir, scan_seconds) - return 1 - - total_cores = cpu_count() - workers = min(args.workers, len(images)) - threads = args.threads or max(1, (total_cores - 1) // workers) - if threads < 1: - logger.error("INVALID_ARGUMENT threads=%d", threads) - return 2 - total_cpu_used = workers * threads - estimated_mem = workers * 2.0 + 2 - - try: - import psutil - - available_gb = psutil.virtual_memory().available / (1024**3) - except ImportError: - available_gb = None - - logger.info( - "PROGRAM_STARTED directory=%s image_count=%d scan_seconds=%.3f workers=%d threads_per_worker=%d total_cores=%d planned_threads=%d reserved_cores=%d stagger_seconds=%.3f estimated_memory_gb=%.1f available_memory_gb=%s", - image_dir, - len(images), - scan_seconds, - workers, - threads, - total_cores, - total_cpu_used, - max(0, total_cores - total_cpu_used), - args.stagger, - estimated_mem, - f"{available_gb:.1f}" if available_gb is not None else "unknown", - ) - - if available_gb is not None and available_gb <= estimated_mem: - logger.warning( - "MEMORY_PRESSURE estimated_memory_gb=%.1f available_memory_gb=%.1f recommendation=reduce_workers", - estimated_mem, - available_gb, - ) - response = input("可用内存可能不足，是否继续？[y/N] ").strip().lower() - if response != "y": - logger.warning("PROGRAM_CANCELLED_BY_USER") - return 0 - - pool_started = time.perf_counter() - results: list[dict] = [] - try: - with Manager() as manager: - log_queue = manager.Queue() - listener_handlers = tuple(logger.handlers) - listener = QueueListener(log_queue, *listener_handlers, respect_handler_level=True) - listener.start() - try: - with Pool( - processes=workers, - initializer=_init_worker, - initargs=(threads, args.stagger, log_queue), - ) as pool: - for completed, result in enumerate( - pool.imap_unordered(_ocr_task, [str(path) for path in images], chunksize=1), - start=1, - ): - results.append(result) - logger.info( - "BATCH_PROGRESS completed=%d total=%d path=%s status=%s image_seconds=%.3f worker_pid=%s", - completed, - len(images), - result["path"], - result["status"], - result["elapsed"], - result.get("worker_pid"), - ) - finally: - listener.stop() - except KeyboardInterrupt: - logger.warning("PROGRAM_INTERRUPTED elapsed_seconds=%.3f", time.perf_counter() - program_started) - return 130 - except Exception as exc: - logger.exception("POOL_FAILED error=%s elapsed_seconds=%.3f", exc, time.perf_counter() - program_started) - return 1 - - pool_seconds = time.perf_counter() - pool_started - completed_results = [result for result in results if result["status"] == "completed"] - failed_results = [result for result in results if result["status"] == "failed"] - worker_metrics = { - result["worker_pid"]: result.get("worker_init", {}) - for result in results - if result.get("worker_pid") is not None - } - worker_model_init_total = sum( - metrics.get("model_init_seconds", 0.0) for metrics in worker_metrics.values() - ) - worker_model_init_average = ( - worker_model_init_total / len(worker_metrics) if worker_metrics else 0.0 - ) - serial_estimate = sum(result["elapsed"] for result in results) - average = serial_estimate / len(results) if results else 0.0 - speedup = serial_estimate / pool_seconds if pool_seconds else 0.0 - program_total = time.perf_counter() - program_started - - for worker_pid, metrics in sorted(worker_metrics.items()): - logger.info( - "WORKER_SUMMARY pid=%s threads=%s stagger_delay_seconds=%s import_seconds=%s model_init_seconds=%s startup_total_seconds=%s priority=%s", - worker_pid, - metrics.get("threads"), - metrics.get("stagger_delay_seconds"), - metrics.get("import_seconds"), - metrics.get("model_init_seconds"), - metrics.get("startup_total_seconds"), - metrics.get("priority"), - ) - - for result in sorted(results, key=lambda item: item["path"]): - if result["status"] == "completed": - logger.info( - "IMAGE_SUMMARY path=%s seconds=%.3f width=%s height=%s layout_boxes=%d parsed_blocks=%d", - result["path"], - result["elapsed"], - result["width"], - result["height"], - result["layout_boxes"], - result["parsed_blocks"], - ) - if not args.no_result: - for index, block in enumerate(result["blocks"], start=1): - logger.info( - "OCR_BLOCK path=%s index=%d label=%s bbox=%s content=%s", - result["path"], - index, - block["label"], - block["bbox"], - block["content"].replace("\r", "").replace("\n", "\\n"), - ) - else: - logger.error("IMAGE_SUMMARY path=%s status=failed error=%s", result["path"], result["error"]) - - logger.info( - "BATCH_SUMMARY image_count=%d completed=%d failed=%d scan_seconds=%.3f pool_seconds=%.3f worker_count=%d worker_model_init_total_seconds=%.3f worker_model_init_average_seconds=%.3f serial_estimate_seconds=%.3f average_image_seconds=%.3f speedup=%.3f program_total_seconds=%.3f workers=%d threads_per_worker=%d log=%s", - len(images), - len(completed_results), - len(failed_results), - scan_seconds, - pool_seconds, - len(worker_metrics), - worker_model_init_total, - worker_model_init_average, - serial_estimate, - average, - speedup, - program_total, - workers, - threads, - log_file.resolve(), - ) - return 0 if not failed_results else 3 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/README.md b/benchmarks/README.md index 893c3ce..7e2e93e 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,5 +1,29 @@ # Benchmark Results -- `gpu/`: GPU Benchmark JSON，由 `gpu/main.py` 生成。 +图片 Benchmark 现在与 OCR 结果保存在同一输出目录： -CPU 当前实测数据记录在根目录 `README.md`。后续可将 CPU 脚本也改为输出同结构 JSON，以进行自动对比。 +```text +outputs/images/<图片名_扩展名>/benchmark.json +``` + +例如： + +```bash +python ocr.py data/images/手写01.png --device cpu --warmup 1 --rounds 3 +``` + +生成： + +```text +outputs/images/手写01_png/benchmark.json +``` + +如需额外复制到指定位置，可使用： + +```bash +python ocr.py data/images/手写01.png \ + --device cpu \ + --benchmark-json benchmarks/手写01-cpu.json +``` + +`benchmarks/cpu/` 与 `benchmarks/gpu/` 仅保留为可选的人工归档目录。 diff --git a/benchmarks/cpu/.gitkeep b/benchmarks/cpu/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/.python-version b/cpu/.python-version similarity index 100% rename from .python-version rename to cpu/.python-version diff --git a/cpu/README.md b/cpu/README.md new file mode 100644 index 0000000..e266502 --- /dev/null +++ b/cpu/README.md @@ -0,0 +1,16 @@ +# CPU 子项目 + +CPU 环境与 GPU 环境完全隔离。通常不直接调用本目录脚本，而是从仓库根目录使用统一入口： + +```bash +python ocr.py data/images/手写01.png --device cpu +python ocr.py data/images/ --device cpu +python ocr.py data/documents/sample.pdf --device cpu --pdf-mode hybrid +python ocr.py verify --device cpu +``` + +安装/更新 CPU 环境： + +```bash +uv sync --project cpu +``` diff --git a/pyproject.toml b/cpu/pyproject.toml similarity index 53% rename from pyproject.toml rename to cpu/pyproject.toml index 1b2673d..4d8983f 100644 --- a/pyproject.toml +++ b/cpu/pyproject.toml @@ -1,7 +1,7 @@ [project] -name = "ocr-vl1-6" -version = "0.1.0" -description = "Add your description here" +name = "ocr-vl1-6-cpu" +version = "0.2.0" +description = "CPU runtime for the unified PaddleOCR-VL-1.6 application" readme = "README.md" requires-python = ">=3.13" dependencies = [ @@ -10,3 +10,8 @@ dependencies = [ "pypdfium2>=5.11.0", "setuptools>=83.0.0", ] + +[dependency-groups] +dev = [ + "pytest>=8.4.0", +] diff --git a/cpu/runner.py b/cpu/runner.py new file mode 100644 index 0000000..4b25695 --- /dev/null +++ b/cpu/runner.py @@ -0,0 +1,14 @@ +"""CPU environment runner used by the root unified launcher.""" + +from pathlib import Path +import sys + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from ocr_app.cli import main + + +if __name__ == "__main__": + raise SystemExit(main(device="cpu")) diff --git a/uv.lock b/cpu/uv.lock similarity index 99% rename from uv.lock rename to cpu/uv.lock index c5d2c0d..c4d6f9b 100644 --- a/uv.lock +++ b/cpu/uv.lock @@ -638,6 +638,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5f/53/fb7122b71361a0d121b669dcf3d31244ef75badbbb724af388948de543e2/imagesize-2.0.0-py2.py3-none-any.whl", hash = "sha256:5667c5bbb57ab3f1fa4bc366f4fbc971db3d5ed011fd2715fd8001f782718d96", size = 9441, upload-time = "2026-03-03T14:18:27.892Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -1026,8 +1035,8 @@ wheels = [ ] [[package]] -name = "ocr-vl1-6" -version = "0.1.0" +name = "ocr-vl1-6-cpu" +version = "0.2.0" source = { virtual = "." } dependencies = [ { name = "paddleocr", extra = ["doc-parser"] }, @@ -1036,6 +1045,11 @@ dependencies = [ { name = "setuptools" }, ] +[package.dev-dependencies] +dev = [ + { name = "pytest" }, +] + [package.metadata] requires-dist = [ { name = "paddleocr", extras = ["doc-parser"], specifier = "==3.7.0" }, @@ -1044,6 +1058,9 @@ requires-dist = [ { name = "setuptools", specifier = ">=83.0.0" }, ] +[package.metadata.requires-dev] +dev = [{ name = "pytest", specifier = ">=8.4.0" }] + [[package]] name = "openai" version = "2.45.0" @@ -1326,6 +1343,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3d/68/1f3066acedf37673694a7141381d8f811ae97f30d34413d236abe7d489f1/pillow-12.3.0-cp315-cp315t-win_arm64.whl", hash = "sha256:06ff022112bc9cbf83b60f8e028d94ad87b60621706487e65f673de61610ab59", size = 2567491, upload-time = "2026-07-01T11:56:23.506Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "premailer" version = "3.10.0" @@ -1608,6 +1634,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f6/d2/42dd53d0a85c27606f316d3aa5d2869c4e8470a5ed6dec30e4a1abe19192/pydantic_core-2.46.4-cp314-cp314t-win_arm64.whl", hash = "sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596", size = 2017325, upload-time = "2026-05-06T13:40:52.723Z" }, ] +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + [[package]] name = "pypdfium2" version = "5.11.0" @@ -1637,6 +1672,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/a7/657503553694c70fba0aa5d213144d80401611c82e0205f43f1ff39f40af/pypdfium2-5.11.0-py3-none-win_arm64.whl", hash = "sha256:897788d71b740752e29875856c369fdb52283f609aecf2355f0473db05dfc1b6", size = 3567726, upload-time = "2026-06-29T11:11:38.689Z" }, ] +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + [[package]] name = "python-bidi" version = "0.6.11" diff --git a/documents/化肥买卖合同　GF—2000—0102.pdf b/data/documents/化肥买卖合同　GF—2000—0102.pdf similarity index 100% rename from documents/化肥买卖合同　GF—2000—0102.pdf rename to data/documents/化肥买卖合同　GF—2000—0102.pdf diff --git a/documents/民用爆破器材买卖合同　GF—2001—0107.pdf b/data/documents/民用爆破器材买卖合同　GF—2001—0107.pdf similarity index 100% rename from documents/民用爆破器材买卖合同　GF—2001—0107.pdf rename to data/documents/民用爆破器材买卖合同　GF—2001—0107.pdf diff --git a/images/名片01.jpg b/data/images/名片01.jpg similarity index 100% rename from images/名片01.jpg rename to data/images/名片01.jpg diff --git a/images/名片02.jpg b/data/images/名片02.jpg similarity index 100% rename from images/名片02.jpg rename to data/images/名片02.jpg diff --git a/images/手写01.png b/data/images/手写01.png similarity index 100% rename from images/手写01.png rename to data/images/手写01.png diff --git a/gpu/README.md b/gpu/README.md index 95fd811..730daae 100644 --- a/gpu/README.md +++ b/gpu/README.md @@ -1,144 +1,64 @@ -# PaddleOCR-VL-1.6 GPU 子项目 +# GPU 子项目 -此目录是与根目录 CPU 版本隔离的 GPU 实验环境。 +GPU 环境与 CPU 环境隔离，所有用户功能通过仓库根目录统一入口调用。 -> **验证状态：未在 NVIDIA GPU 上实测。** 当前开发机器只有集成显卡，无法运行 CUDA。代码仅完成静态检查；最终安装、兼容性和性能必须在目标 NVIDIA GPU 机器上验证。 - -## 为什么独立环境 - -`paddlepaddle` 与 `paddlepaddle-gpu` 都提供 `paddle` 模块，不能安全共用同一个虚拟环境。本目录具有独立的： - -- `pyproject.toml` -- `.python-version`（Python 3.11） -- `.venv`（执行安装后生成） -- `uv.lock`（在目标 GPU 机器安装后生成） - -根目录 CPU 环境不会被修改。 - -## 前置条件 - -1. NVIDIA CUDA GPU（Intel/AMD 集成显卡不能运行 Paddle CUDA 版本） -2. 兼容的 NVIDIA 驱动 -3. Python 3.11 -4. uv -5. 根据 PaddlePaddle 官方兼容表选择 CUDA Wheel - -先检查目标机器： - -```bash -nvidia-smi -``` - -`nvidia-smi` 显示的 CUDA Version 是驱动支持上限，不等同于本机安装的 CUDA Toolkit，也不能单独用于判断 Wheel 版本。 +> 当前开发机器没有 NVIDIA 独立显卡，GPU 实机功能尚未验证。 ## 安装 -在仓库根目录运行： +依据 PaddlePaddle 官方兼容表选择 CUDA Wheel： ```bash -# 只查看将执行的命令，不安装 python gpu/setup_env.py --cuda cu118 --dry-run - -# CUDA 11.8 Wheel python gpu/setup_env.py --cuda cu118 +``` -# 或 CUDA 12.6 Wheel +或： + +```bash python gpu/setup_env.py --cuda cu126 ``` -安装脚本将在 `gpu/.venv` 创建独立环境。若官方 Wheel 支持范围发生变化，请同步更新 `gpu/pyproject.toml` 与 `gpu/setup_env.py`。 +安装脚本会： -## 验证环境 +1. 创建 `gpu/.venv` +2. 使用 PaddlePaddle CUDA 专用索引安装依赖 +3. 成功后创建 `gpu/.gpu-ready` + +无 `nvidia-smi` 时默认拒绝安装。`--allow-no-gpu` 仅用于准备环境或 CI，不代表环境可以推理。 + +## 统一入口 ```bash -uv run --project gpu python gpu/verify_env.py +python ocr.py verify --device gpu +python ocr.py data/images/手写01.png --device gpu --warmup 1 --rounds 3 +python ocr.py data/images/ --device gpu +python ocr.py data/documents/sample.pdf --device gpu --pdf-mode hybrid ``` -该脚本会检查： +统一入口只调用已经安装完成的 `gpu/.venv`，不会从默认 PyPI 重新解析 `paddlepaddle-gpu`，也不会自动回退到 CPU。 -- PaddlePaddle 是否为 CUDA 构建 -- CUDA GPU 数量及名称 -- `gpu:0` 是否能完成矩阵乘法 +## PDF 混合模式 -任何检查失败都会以非零状态退出，不会自动回退到 CPU。 - -## 单图 Benchmark +`--pdf-mode hybrid` 会先读取 PDF 文本层。只有需要 OCR 的页面才初始化 GPU 模型，因此纯电子 PDF 不会创建 CUDA 模型或占用大块显存。 ```bash -uv run --project gpu python gpu/main.py -``` - -常用参数： - -```bash -uv run --project gpu python gpu/main.py \ - --image images/手写01.png \ +python ocr.py data/documents/sample.pdf \ + --device gpu \ --device-id 0 \ - --warmup 1 \ - --rounds 3 + --pdf-mode hybrid ``` -Windows PowerShell 可写为单行，或使用反引号续行。 +## 当前限制 -Benchmark 会记录： +- 仅单 GPU +- 批量图片使用单模型串行处理 +- 未启用 TensorRT/FP16/BF16 +- 未验证 PaddleOCR-VL-1.6 在目标显卡上的显存需求 +- 未实现多 GPU 调度 -- GPU 型号和设备编号 -- Python/PaddlePaddle 版本 -- 模型初始化耗时 -- 预热和正式推理轮数 -- min/max/mean/median/stdev -- 可获取时的 CUDA 显存统计 -- 图片尺寸和文本块数量 - -结果写入： - -```text -benchmarks/gpu/gpu-benchmark-YYYYMMDD-HHMMSS.json -logs/single/<图片名>-gpuN-YYYYMMDD-HHMMSS.log -``` - -日志记录 CUDA 配置、PaddleOCR 导入、模型初始化、每轮预热/推理、显存统计和程序总用时。可用 `--log-file` 指定路径，使用 `--verbose` 输出详细异常。 - -## PDF OCR - -GPU PDF 入口复用仓库根目录 `pdf_ocr_core.py`，按页渲染、逐页保存并支持断点续传： +目标 GPU 机器安装后，先执行： ```bash -uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf \ - --device-id 0 \ - --pages "1-10" \ - --dpi 144 +python ocr.py verify --device gpu ``` - -常用选项： - -```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 实机验证。 - -PDF 日志默认写入： - -```text -logs/pdf/-gpuN-YYYYMMDD-HHMMSS.log -``` - -日志与 `manifest.json` 会记录每页渲染、OCR、结果导出、状态保存、任务总用时和程序总用时。 - -## 当前范围 - -当前实现单 GPU、单图 Benchmark 和单 GPU PDF 逐页 OCR。暂未实现 GPU 多进程批处理，原因是： - -- 同一 GPU 上启动多个模型实例会重复占用显存 -- 多进程通常不会线性提升单卡吞吐 -- 容易引发显存不足和 CUDA 上下文争抢 - -后续应优先评估模型/pipeline 原生批处理能力，再决定是否增加多 GPU 或任务队列。 diff --git a/gpu/main.py b/gpu/main.py deleted file mode 100644 index 9fbec3f..0000000 --- a/gpu/main.py +++ /dev/null @@ -1,291 +0,0 @@ -"""PaddleOCR-VL-1.6 GPU 单图推理与 Benchmark。""" - -import argparse -import json -import platform -import statistics -import sys -import time -from datetime import datetime -from pathlib import Path -from typing import Any - -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 ocr_logging import default_log_path, setup_run_logger -DEFAULT_IMAGE = PROJECT_ROOT / "images" / "手写01.png" -DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "benchmarks" / "gpu" - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="PaddleOCR-VL-1.6 GPU Benchmark") - parser.add_argument("--image", type=Path, default=DEFAULT_IMAGE, help="待识别图片") - parser.add_argument("--device-id", type=int, default=0, help="CUDA GPU 编号") - parser.add_argument("--warmup", type=int, default=1, help="预热轮数") - parser.add_argument("--rounds", type=int, default=3, help="正式测试轮数") - parser.add_argument( - "--output-dir", - type=Path, - default=DEFAULT_OUTPUT_DIR, - help="Benchmark JSON 输出目录", - ) - parser.add_argument("--no-result", action="store_true", help="不在控制台输出 OCR 文本") - parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") - parser.add_argument("--verbose", action="store_true", help="输出详细日志") - return parser.parse_args() - - -def validate_args(args: argparse.Namespace) -> None: - args.image = args.image.expanduser().resolve() - args.output_dir = args.output_dir.expanduser().resolve() - if not args.image.is_file(): - raise ValueError(f"图片不存在: {args.image}") - if args.device_id < 0: - raise ValueError("--device-id 不能小于 0") - if args.warmup < 0: - raise ValueError("--warmup 不能小于 0") - if args.rounds < 1: - raise ValueError("--rounds 必须大于等于 1") - - -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 支持。请确认安装的是 paddlepaddle-gpu，" - "且正在使用 gpu/.venv。" - ) - - 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 >= device_count: - raise RuntimeError(f"GPU {device_id} 不存在，当前仅检测到 {device_count} 个 CUDA 设备。") - - device = f"gpu:{device_id}" - try: - paddle.set_device(device) - paddle.device.cuda.synchronize(device_id) - except Exception as exc: - raise RuntimeError(f"无法启用 {device}: {exc}") from exc - - try: - device_name = paddle.device.cuda.get_device_name(device_id) - except Exception: - device_name = "unknown" - - return paddle, device, device_name - - -def synchronize(paddle: Any, device_id: int) -> None: - paddle.device.cuda.synchronize(device_id) - - -def read_gpu_memory(paddle: Any, device_id: int) -> dict[str, float | None]: - stats: dict[str, float | None] = { - "allocated_mb": None, - "reserved_mb": None, - "max_allocated_mb": None, - "max_reserved_mb": None, - } - functions = { - "allocated_mb": "memory_allocated", - "reserved_mb": "memory_reserved", - "max_allocated_mb": "max_memory_allocated", - "max_reserved_mb": "max_memory_reserved", - } - for key, function_name in functions.items(): - function = getattr(paddle.device.cuda, function_name, None) - if function is None: - continue - try: - stats[key] = round(float(function(device_id)) / (1024**2), 2) - except Exception: - pass - return stats - - -def result_summary(result: list[Any]) -> dict[str, Any]: - first = result[0] - blocks = first["parsing_res_list"] - return { - "width": first["width"], - "height": first["height"], - "layout_boxes": len(first["layout_det_res"]["boxes"]), - "parsed_blocks": len(blocks), - "non_empty_blocks": sum(bool(block.content.strip()) for block in blocks), - } - - -def print_ocr_result(result: list[Any]) -> None: - print("\n[OCR Result]") - for item in result: - for block in item["parsing_res_list"]: - if block.content.strip(): - print(f"[{block.label}] {block.bbox}") - print(block.content) - print() - - -def main() -> int: - program_started = time.perf_counter() - args = parse_args() - log_file = args.log_file or default_log_path( - PROJECT_ROOT, - "single", - args.image.stem, - device=f"gpu{args.device_id}", - ) - logger = setup_run_logger("ocr.single.gpu", log_file, verbose=args.verbose) - logger.info( - "PROGRAM_STARTED image=%s device_id=%d warmup=%d rounds=%d output_dir=%s", - args.image, - args.device_id, - args.warmup, - args.rounds, - args.output_dir, - ) - try: - validate_args(args) - cuda_started = time.perf_counter() - paddle, device, device_name = configure_cuda(args.device_id) - cuda_setup_seconds = time.perf_counter() - cuda_started - except (ValueError, RuntimeError) as exc: - logger.error("VALIDATION_OR_CUDA_FAILED type=%s error=%s", type(exc).__name__, exc, exc_info=args.verbose) - return 1 - - import_started = time.perf_counter() - from paddleocr import PaddleOCRVL - import_seconds = time.perf_counter() - import_started - logger.info( - "RUNTIME_READY cuda_setup_seconds=%.3f import_seconds=%.3f device=%s device_name=%s paddle_version=%s image_size_bytes=%d", - cuda_setup_seconds, - import_seconds, - device, - device_name, - paddle.__version__, - args.image.stat().st_size, - ) - - logger.info("MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=%s", device) - synchronize(paddle, args.device_id) - init_started = time.perf_counter() - pipeline = PaddleOCRVL(pipeline_version="v1.6", device=device) - synchronize(paddle, args.device_id) - init_seconds = time.perf_counter() - init_started - logger.info("MODEL_INITIALIZED seconds=%.3f", init_seconds) - - result = None - warmup_times: list[float] = [] - for index in range(args.warmup): - started = time.perf_counter() - result = pipeline.predict(str(args.image)) - synchronize(paddle, args.device_id) - elapsed = time.perf_counter() - started - warmup_times.append(elapsed) - logger.info("WARMUP_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.warmup, elapsed) - - inference_times: list[float] = [] - for index in range(args.rounds): - synchronize(paddle, args.device_id) - started = time.perf_counter() - result = pipeline.predict(str(args.image)) - synchronize(paddle, args.device_id) - elapsed = time.perf_counter() - started - inference_times.append(elapsed) - logger.info("INFERENCE_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.rounds, elapsed) - - if result is None: - logger.error("EMPTY_RESULT") - return 2 - - summary = result_summary(result) - benchmark = { - "status": "completed", - "timestamp": datetime.now().astimezone().isoformat(), - "platform": platform.platform(), - "python_version": platform.python_version(), - "paddle_version": paddle.__version__, - "pipeline_version": "v1.6", - "device": device, - "device_name": device_name, - "image_path": str(args.image), - "image": summary, - "warmup_rounds": args.warmup, - "benchmark_rounds": args.rounds, - "cuda_setup_seconds": round(cuda_setup_seconds, 3), - "runtime_import_seconds": round(import_seconds, 3), - "model_init_seconds": round(init_seconds, 3), - "warmup_seconds": [round(value, 3) for value in warmup_times], - "inference_seconds": { - "all": [round(value, 3) for value in inference_times], - "min": round(min(inference_times), 3), - "max": round(max(inference_times), 3), - "mean": round(statistics.fmean(inference_times), 3), - "median": round(statistics.median(inference_times), 3), - "stdev": round(statistics.pstdev(inference_times), 3), - }, - "gpu_memory": read_gpu_memory(paddle, args.device_id), - "program_total_seconds": round(time.perf_counter() - program_started, 3), - "log_file": str(log_file.resolve()), - } - - args.output_dir.mkdir(parents=True, exist_ok=True) - timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") - output_path = args.output_dir / f"gpu-benchmark-{timestamp}.json" - output_path.write_text(json.dumps(benchmark, ensure_ascii=False, indent=2), encoding="utf-8") - - logger.info( - "RESULT_SUMMARY width=%d height=%d layout_boxes=%d parsed_blocks=%d non_empty_blocks=%d gpu_memory=%s", - summary["width"], - summary["height"], - summary["layout_boxes"], - summary["parsed_blocks"], - summary["non_empty_blocks"], - benchmark["gpu_memory"], - ) - logger.info( - "BENCHMARK_SUMMARY cuda_setup_seconds=%.3f import_seconds=%.3f model_init_seconds=%.3f warmup_total_seconds=%.3f inference_total_seconds=%.3f inference_min_seconds=%.3f inference_max_seconds=%.3f inference_mean_seconds=%.3f inference_median_seconds=%.3f inference_stdev_seconds=%.3f program_total_seconds=%.3f result_json=%s log=%s", - cuda_setup_seconds, - import_seconds, - init_seconds, - sum(warmup_times), - sum(inference_times), - min(inference_times), - max(inference_times), - statistics.fmean(inference_times), - statistics.median(inference_times), - statistics.pstdev(inference_times), - time.perf_counter() - program_started, - output_path, - log_file.resolve(), - ) - - if not args.no_result: - for index, block in enumerate(result[0]["parsing_res_list"], start=1): - if block.content.strip(): - logger.info( - "OCR_BLOCK index=%d label=%s bbox=%s content=%s", - index, - block.label, - block.bbox, - block.content.replace("\r", "").replace("\n", "\\n"), - ) - - logger.info("PROGRAM_COMPLETED") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/gpu/pdf_ocr.py b/gpu/pdf_ocr.py deleted file mode 100644 index 3340faf..0000000 --- a/gpu/pdf_ocr.py +++ /dev/null @@ -1,214 +0,0 @@ -"""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 ocr_logging import default_log_path, setup_run_logger -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 最大输入像素参数") - parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") - parser.add_argument("--verbose", action="store_true", help="输出详细日志") - 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: - program_started = time.perf_counter() - args = parse_args() - log_file = args.log_file or default_log_path( - PROJECT_ROOT, - "pdf", - args.pdf.stem, - device=f"gpu{args.device_id}", - ) - logger = setup_run_logger("ocr.pdf.gpu", log_file, verbose=args.verbose) - logger.info( - "PROGRAM_STARTED input=%s output=%s pages=%s dpi=%d device_id=%d resume=%s overwrite=%s keep_rendered=%s fail_fast=%s", - args.pdf, - args.output, - args.pages or "all", - args.dpi, - args.device_id, - args.resume, - args.overwrite, - args.keep_rendered, - args.fail_fast, - ) - try: - preflight_started = time.perf_counter() - 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, - ) - preflight_seconds = time.perf_counter() - preflight_started - cuda_started = time.perf_counter() - paddle, device, device_name = configure_cuda(args.device_id) - cuda_setup_seconds = time.perf_counter() - cuda_started - except Exception as exc: - logger.error("PREFLIGHT_OR_CUDA_FAILED type=%s error=%s", type(exc).__name__, exc, exc_info=args.verbose) - return 1 - - import_started = time.perf_counter() - from paddleocr import PaddleOCRVL - import_seconds = time.perf_counter() - import_started - - logger.info( - "PREFLIGHT_COMPLETED seconds=%.3f page_count=%d selected_pages=%d document_dir=%s", - preflight_seconds, - preflight["page_count"], - len(preflight["selected_pages"]), - preflight["document_dir"], - ) - logger.info( - "RUNTIME_READY cuda_setup_seconds=%.3f import_seconds=%.3f device=%s device_name=%s paddle_version=%s", - cuda_setup_seconds, - import_seconds, - device, - device_name, - paddle.__version__, - ) - logger.info("MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=%s", device) - 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 - logger.info("MODEL_INITIALIZED seconds=%.3f pipeline_version=v1.6 device=%s", init_seconds, device) - - 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", - "preflight_seconds": round(preflight_seconds, 3), - "cuda_setup_seconds": round(cuda_setup_seconds, 3), - "runtime_import_seconds": round(import_seconds, 3), - "log_file": str(log_file.resolve()), - } - - 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), - logger=logger, - ) - except KeyboardInterrupt: - logger.warning( - "PROGRAM_INTERRUPTED total_seconds=%.3f resume_hint=--resume", - time.perf_counter() - program_started, - ) - return 130 - except Exception as exc: - logger.exception( - "PROGRAM_FAILED type=%s error=%s total_seconds=%.3f", - type(exc).__name__, - exc, - time.perf_counter() - program_started, - ) - return 1 - - program_total = time.perf_counter() - program_started - timing = summary.get("timing", {}) - logger.info( - "PROGRAM_COMPLETED status=%s completed_pages=%d selected_pages=%d failed_pages=%s model_init_seconds=%.3f pdf_task_seconds=%.3f program_total_seconds=%.3f output=%s log=%s", - summary["status"], - summary["completed_pages"], - summary["selected_pages"], - summary["failed_pages"], - init_seconds, - timing.get("task_total_seconds", 0.0), - program_total, - summary["document_dir"], - log_file.resolve(), - ) - 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 97ae644..5f58484 100644 --- a/gpu/pyproject.toml +++ b/gpu/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ocr-vl1-6-gpu" -version = "0.1.0" -description = "GPU benchmark for PaddleOCR-VL-1.6" +version = "0.2.0" +description = "GPU runtime for the unified PaddleOCR-VL-1.6 application" readme = "README.md" requires-python = ">=3.11,<3.13" dependencies = [ diff --git a/gpu/runner.py b/gpu/runner.py new file mode 100644 index 0000000..c7c0c3e --- /dev/null +++ b/gpu/runner.py @@ -0,0 +1,14 @@ +"""GPU environment runner used by the root unified launcher.""" + +from pathlib import Path +import sys + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from ocr_app.cli import main + + +if __name__ == "__main__": + raise SystemExit(main(device="gpu")) diff --git a/gpu/setup_env.py b/gpu/setup_env.py index b47895d..f3c0801 100644 --- a/gpu/setup_env.py +++ b/gpu/setup_env.py @@ -67,6 +67,10 @@ def main() -> int: if args.dry_run: return 0 + ready_marker = project_dir / ".gpu-ready" + if ready_marker.exists(): + ready_marker.unlink() + completed = subprocess.run(command, check=False) if completed.returncode != 0: print( @@ -75,8 +79,12 @@ def main() -> int: ) return completed.returncode - print("\n[OK] GPU 子项目环境已创建。下一步运行：") - print(f' uv run --project "{project_dir}" python "{project_dir / "verify_env.py"}"') + ready_marker.write_text( + f"cuda={args.cuda}\nindex={index_url}\n", + encoding="utf-8", + ) + print("\n[OK] GPU 子项目环境已创建。下一步从仓库根目录运行：") + print(" python ocr.py verify --device gpu") return 0 diff --git a/gpu/verify_env.py b/gpu/verify_env.py deleted file mode 100644 index b819e7d..0000000 --- a/gpu/verify_env.py +++ /dev/null @@ -1,53 +0,0 @@ -"""只检查 GPU Paddle 环境，不加载 OCR 模型。""" - -import sys - - -def main() -> int: - try: - import paddle - except ImportError: - print("[ERROR] 未安装 PaddlePaddle GPU。请先运行 setup_env.py。") - return 1 - - print(f"Python: {sys.version.split()[0]}") - print(f"PaddlePaddle: {paddle.__version__}") - print(f"CUDA build: {paddle.is_compiled_with_cuda()}") - - if not paddle.is_compiled_with_cuda(): - print("[ERROR] 当前安装的 PaddlePaddle 不是 CUDA 版本。") - return 2 - - try: - device_count = paddle.device.cuda.device_count() - except Exception as exc: - print(f"[ERROR] 无法查询 CUDA 设备: {exc}") - return 3 - - print(f"CUDA device count: {device_count}") - if device_count < 1: - print("[ERROR] 未检测到可用的 NVIDIA CUDA GPU。") - return 4 - - for device_id in range(device_count): - try: - name = paddle.device.cuda.get_device_name(device_id) - except Exception: - name = "unknown" - print(f"GPU {device_id}: {name}") - - try: - paddle.set_device("gpu:0") - tensor = paddle.ones([1024, 1024], dtype="float32") - result = paddle.matmul(tensor, tensor) - paddle.device.cuda.synchronize() - print(f"CUDA smoke test: OK, result shape={list(result.shape)}") - except Exception as exc: - print(f"[ERROR] CUDA 计算测试失败: {exc}") - return 5 - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/logs/名片01.log b/logs/legacy/名片01.log similarity index 100% rename from logs/名片01.log rename to logs/legacy/名片01.log diff --git a/logs/名片02.log b/logs/legacy/名片02.log similarity index 100% rename from logs/名片02.log rename to logs/legacy/名片02.log diff --git a/logs/手写01.log b/logs/legacy/手写01.log similarity index 100% rename from logs/手写01.log rename to logs/legacy/手写01.log diff --git a/logs/批量识别.log b/logs/legacy/批量识别.log similarity index 100% rename from logs/批量识别.log rename to logs/legacy/批量识别.log diff --git a/main.py b/main.py deleted file mode 100644 index a8f0a8c..0000000 --- a/main.py +++ /dev/null @@ -1,140 +0,0 @@ -"""CPU single-image OCR benchmark with structured timing logs.""" - -from __future__ import annotations - -import argparse -import os -import statistics -import time -from pathlib import Path - -from ocr_logging import default_log_path, setup_run_logger - -PROJECT_ROOT = Path(__file__).resolve().parent -DEFAULT_IMAGE = PROJECT_ROOT / "images" / "名片02.jpg" - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="PaddleOCR-VL-1.6 CPU 单图 Benchmark", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument("image", nargs="?", type=Path, default=DEFAULT_IMAGE, help="输入图片") - parser.add_argument("--threads", type=int, default=None, help="Paddle CPU 线程数") - parser.add_argument("--warmup", type=int, default=0, help="预热轮数") - parser.add_argument("--rounds", type=int, default=1, help="正式测试轮数") - parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") - parser.add_argument("--verbose", action="store_true", help="输出详细日志") - parser.add_argument("--no-result", action="store_true", help="不输出识别文本") - return parser.parse_args() - - -def main() -> int: - program_started = time.perf_counter() - args = parse_args() - image_path = args.image.expanduser().resolve() - log_file = args.log_file or default_log_path(PROJECT_ROOT, "single", image_path.stem, device="cpu") - logger = setup_run_logger("ocr.single.cpu", log_file, verbose=args.verbose) - - if not image_path.is_file(): - logger.error("INPUT_NOT_FOUND path=%s", image_path) - return 1 - if args.warmup < 0 or args.rounds < 1: - logger.error("INVALID_ARGUMENT warmup=%d rounds=%d", args.warmup, args.rounds) - return 2 - - total_cores = os.cpu_count() or 4 - threads = args.threads or int(os.environ.get("PADDLE_THREADS", total_cores)) - if threads < 1: - logger.error("INVALID_ARGUMENT threads=%d", threads) - return 2 - - logger.info( - "PROGRAM_STARTED image=%s size_bytes=%d threads=%d total_cores=%d warmup=%d rounds=%d", - image_path, - image_path.stat().st_size, - threads, - total_cores, - args.warmup, - args.rounds, - ) - - import_started = time.perf_counter() - from paddle import core - from paddleocr import PaddleOCRVL - import_seconds = time.perf_counter() - import_started - core.set_num_threads(threads) - logger.info("RUNTIME_READY import_seconds=%.3f threads=%d", import_seconds, threads) - - logger.info("MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=cpu") - init_started = time.perf_counter() - pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu") - init_seconds = time.perf_counter() - init_started - logger.info("MODEL_INITIALIZED seconds=%.3f", init_seconds) - - warmup_times: list[float] = [] - for index in range(args.warmup): - started = time.perf_counter() - pipeline.predict(str(image_path)) - elapsed = time.perf_counter() - started - warmup_times.append(elapsed) - logger.info("WARMUP_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.warmup, elapsed) - - result = None - inference_times: list[float] = [] - for index in range(args.rounds): - started = time.perf_counter() - result = pipeline.predict(str(image_path)) - elapsed = time.perf_counter() - started - inference_times.append(elapsed) - logger.info("INFERENCE_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.rounds, elapsed) - - if not result: - logger.error("EMPTY_RESULT") - return 3 - - first = result[0] - layout_boxes = len(first["layout_det_res"]["boxes"]) - parsed_blocks = len(first["parsing_res_list"]) - non_empty_blocks = sum(bool(block.content.strip()) for block in first["parsing_res_list"]) - inference_mean = statistics.fmean(inference_times) - inference_stdev = statistics.pstdev(inference_times) - program_total = time.perf_counter() - program_started - - logger.info( - "RESULT_SUMMARY width=%s height=%s layout_boxes=%d parsed_blocks=%d non_empty_blocks=%d", - first.get("width"), - first.get("height"), - layout_boxes, - parsed_blocks, - non_empty_blocks, - ) - logger.info( - "BENCHMARK_SUMMARY model_init_seconds=%.3f warmup_total_seconds=%.3f inference_total_seconds=%.3f inference_min_seconds=%.3f inference_max_seconds=%.3f inference_mean_seconds=%.3f inference_median_seconds=%.3f inference_stdev_seconds=%.3f program_total_seconds=%.3f", - init_seconds, - sum(warmup_times), - sum(inference_times), - min(inference_times), - max(inference_times), - inference_mean, - statistics.median(inference_times), - inference_stdev, - program_total, - ) - - if not args.no_result: - for index, block in enumerate(first["parsing_res_list"], start=1): - logger.info( - "OCR_BLOCK index=%d label=%s bbox=%s content=%s", - index, - block.label, - block.bbox, - block.content.replace("\r", "").replace("\n", "\\n"), - ) - - logger.info("PROGRAM_COMPLETED log=%s", log_file.resolve()) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/ocr.py b/ocr.py new file mode 100644 index 0000000..8aeb598 --- /dev/null +++ b/ocr.py @@ -0,0 +1,78 @@ +"""Unified launcher. + +Examples: + python ocr.py data/images/手写01.png --device cpu + python ocr.py data/documents/sample.pdf --device cpu --pdf-mode hybrid + python ocr.py data/ --recursive --device cpu + python ocr.py verify --device gpu + +The launcher deliberately executes the selected isolated uv project, so CPU +and GPU Paddle packages never share the same virtual environment. +""" + +from __future__ import annotations + +import shutil +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parent + + +def _requested_device(argv: list[str]) -> str: + for index, value in enumerate(argv): + if value == "--device": + if index + 1 >= len(argv): + raise SystemExit("--device 需要 cpu 或 gpu") + return argv[index + 1].lower() + if value.startswith("--device="): + return value.split("=", 1)[1].lower() + return "cpu" + + +def main() -> int: + try: + device = _requested_device(sys.argv[1:]) + except SystemExit as exc: + print(exc, file=sys.stderr) + return 2 + if device not in {"cpu", "gpu"}: + print(f"不支持的设备: {device}，可选 cpu/gpu", file=sys.stderr) + return 2 + + project = ROOT / device + runner = project / "runner.py" + if sys.platform == "win32": + environment_python = project / ".venv" / "Scripts" / "python.exe" + else: + environment_python = project / ".venv" / "bin" / "python" + + if device == "gpu": + ready_marker = project / ".gpu-ready" + if not ready_marker.is_file() or not environment_python.is_file(): + print( + "GPU 环境尚未安装完成。请根据目标 CUDA 版本运行：\n" + " python gpu/setup_env.py --cuda cu118\n" + "或：\n" + " python gpu/setup_env.py --cuda cu126", + file=sys.stderr, + ) + return 1 + command = [str(environment_python), str(runner), *sys.argv[1:]] + return subprocess.run(command, cwd=ROOT).returncode + + if environment_python.is_file(): + command = [str(environment_python), str(runner), *sys.argv[1:]] + return subprocess.run(command, cwd=ROOT).returncode + + uv = shutil.which("uv") + if not uv: + print("未找到 CPU 虚拟环境和 uv，请先安装 uv。", file=sys.stderr) + return 1 + command = [uv, "run", "--project", str(project), "python", str(runner), *sys.argv[1:]] + return subprocess.run(command, cwd=ROOT).returncode + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ocr_app/__init__.py b/ocr_app/__init__.py new file mode 100644 index 0000000..95cb432 --- /dev/null +++ b/ocr_app/__init__.py @@ -0,0 +1,3 @@ +"""Shared PaddleOCR-VL application package.""" + +__version__ = "0.2.0" diff --git a/ocr_app/cli.py b/ocr_app/cli.py new file mode 100644 index 0000000..e13a7ef --- /dev/null +++ b/ocr_app/cli.py @@ -0,0 +1,124 @@ +"""Path-first unified CLI: one file or one directory, routed by suffix.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +from .commands import run_input, run_verify +from .logging_utils import default_log_path, setup_run_logger +from .runtime import PipelineProvider, RuntimeConfig + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +LEGACY_COMMANDS = {"image", "pdf", "batch"} + + +def _add_device_options(parser: argparse.ArgumentParser, default: str | None) -> None: + parser.add_argument("--device", choices=("cpu", "gpu"), default=default or "cpu", help="运行设备") + parser.add_argument("--device-id", type=int, default=0, help="GPU 编号") + parser.add_argument("--threads", type=int, default=None, help="CPU 线程数") + parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") + parser.add_argument("--verbose", action="store_true", help="输出详细日志") + + +def build_input_parser(device_override: str | None = None) -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="ocr.py", + description="自动按文件后缀处理图片/PDF；输入目录时批量使用同一路由逻辑", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("input", type=Path, help="图片、PDF 或目录") + parser.add_argument("--recursive", action="store_true", help="目录模式递归扫描子目录") + parser.add_argument("--output", type=Path, default=PROJECT_ROOT / "outputs", help="PDF 输出根目录") + parser.add_argument("--fail-fast", action="store_true", help="单文件失败后立即停止目录任务") + + # Image options + parser.add_argument("--warmup", type=int, default=0, help="首张图片预热轮数") + parser.add_argument("--rounds", type=int, default=1, help="每张图片推理轮数") + parser.add_argument("--benchmark-json", type=Path, default=None, help="单图片 Benchmark JSON 路径") + parser.add_argument("--no-result", action="store_true", help="不记录图片 OCR 文本块") + + # PDF options + parser.add_argument( + "--pdf-mode", + "--mode", + dest="pdf_mode", + choices=("hybrid", "text", "ocr"), + default="hybrid", + help="PDF 处理模式", + ) + parser.add_argument("--pages", help="PDF 页码范围，例如 1-5,8,10-") + parser.add_argument("--dpi", type=int, default=144, help="PDF OCR 页面渲染 DPI") + parser.add_argument("--password", help="PDF 密码") + parser.add_argument("--resume", action="store_true", help="PDF 断点续传") + parser.add_argument("--overwrite", action="store_true", help="覆盖已有 PDF 输出") + parser.add_argument("--keep-rendered", action="store_true", help="保留 OCR 页面 PNG") + parser.add_argument("--max-new-tokens", type=int, default=None, help="VLM 最大生成 token") + parser.add_argument("--min-pixels", type=int, default=None, help="VLM 最小输入像素") + parser.add_argument("--max-pixels", type=int, default=None, help="VLM 最大输入像素") + parser.add_argument("--text-min-chars", type=int, default=50, help="有效文本最小字符数") + parser.add_argument("--text-min-printable-ratio", type=float, default=0.85, help="可打印字符比例阈值") + parser.add_argument("--text-min-content-ratio", type=float, default=0.60, help="字母/数字/CJK 比例阈值") + parser.add_argument("--text-max-replacement-ratio", type=float, default=0.02, help="替换字符最大比例") + parser.add_argument("--text-min-density", type=float, default=25.0, help="文本密度阈值") + _add_device_options(parser, device_override) + return parser + + +def build_verify_parser(device_override: str | None = None) -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="ocr.py verify", + description="验证 CPU/GPU Paddle 环境", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + _add_device_options(parser, device_override) + return parser + + +def normalize_argv(argv: list[str]) -> tuple[str, list[str]]: + """Keep old image/pdf/batch prefixes as compatibility aliases.""" + if argv and argv[0] == "verify": + return "verify", argv[1:] + if argv and argv[0] in LEGACY_COMMANDS: + return "input", argv[1:] + return "input", argv + + +def main(device: str | None = None, argv: list[str] | None = None) -> int: + raw_argv = list(sys.argv[1:] if argv is None else argv) + command, normalized = normalize_argv(raw_argv) + parser = build_verify_parser(device) if command == "verify" else build_input_parser(device) + args = parser.parse_args(normalized) + if device is not None: + args.device = device + + if command == "verify": + stem = "verify" + category = "verify" + else: + stem = args.input.stem or args.input.name or "input" + category = "input" + log_file = args.log_file or default_log_path(PROJECT_ROOT, category, stem, device=args.device) + logger = setup_run_logger(f"ocr.{category}.{args.device}", log_file, verbose=args.verbose) + logger.info( + "COMMAND_STARTED command=%s input=%s device=%s log=%s", + command, + getattr(args, "input", None), + args.device, + log_file, + ) + + provider = PipelineProvider( + RuntimeConfig(device=args.device, threads=args.threads, device_id=args.device_id), + logger, + ) + try: + if command == "verify": + return run_verify(args, provider, logger, PROJECT_ROOT) + if args.warmup < 0 or args.rounds < 1: + raise ValueError("--warmup 必须 >= 0，--rounds 必须 >= 1") + return run_input(args, provider, logger, PROJECT_ROOT) + except Exception as exc: + logger.exception("COMMAND_FAILED command=%s error=%s", command, exc) + return 1 diff --git a/ocr_app/commands.py b/ocr_app/commands.py new file mode 100644 index 0000000..6a74f86 --- /dev/null +++ b/ocr_app/commands.py @@ -0,0 +1,497 @@ +"""Unified suffix-based routing for files and directories.""" + +from __future__ import annotations + +import logging +import statistics +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any + +from .output import ( + atomic_write_json, + image_output_directory, + pdf_output_root, + safe_stem, + save_image_ocr_outputs, +) +from .pdf import preflight_pdf, process_pdf +from .pdf_text import TextLayerPolicy +from .runtime import PipelineProvider + +IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".webp"} +PDF_EXTENSIONS = {".pdf"} +SUPPORTED_EXTENSIONS = IMAGE_EXTENSIONS | PDF_EXTENSIONS + + +@dataclass +class FileProcessResult: + path: Path + kind: str + status: str + seconds: float + details: dict[str, Any] + exit_code: int = 0 + + +def detect_input_kind(path: Path) -> str: + if path.is_dir(): + return "directory" + suffix = path.suffix.lower() + if suffix in IMAGE_EXTENSIONS: + return "image" + if suffix in PDF_EXTENSIONS: + return "pdf" + return "unsupported" + + +def discover_supported_files(directory: Path, *, recursive: bool, output_root: Path) -> list[Path]: + iterator = directory.rglob("*") if recursive else directory.glob("*") + output_root = output_root.expanduser().resolve() + files: list[Path] = [] + for path in iterator: + if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS: + continue + resolved = path.resolve() + try: + resolved.relative_to(output_root) + except ValueError: + files.append(resolved) + return sorted(files, key=lambda value: str(value).casefold()) + + +def _result_summary(result: list[Any]) -> dict[str, Any]: + first = result[0] + blocks = first["parsing_res_list"] + return { + "width": first.get("width"), + "height": first.get("height"), + "layout_boxes": len(first.get("layout_det_res", {}).get("boxes", [])), + "parsed_blocks": len(blocks), + "non_empty_blocks": sum(bool(block.content.strip()) for block in blocks), + } + + +def process_image_file( + path: Path, + *, + args, + provider: PipelineProvider, + logger: logging.Logger, + project_root: Path, + run_warmup: bool, + batch_root: Path | None, +) -> FileProcessResult: + file_started = time.perf_counter() + logger.info("FILE_ROUTED path=%s kind=image", path) + pipeline = provider.get() + predict_kwargs = { + key: value + for key, value in { + "max_new_tokens": getattr(args, "max_new_tokens", None), + "min_pixels": getattr(args, "min_pixels", None), + "max_pixels": getattr(args, "max_pixels", None), + }.items() + if value is not None + } + + warmup_times: list[float] = [] + if run_warmup: + for index in range(args.warmup): + started = time.perf_counter() + pipeline.predict(str(path), **predict_kwargs) + provider.synchronize() + elapsed = time.perf_counter() - started + warmup_times.append(elapsed) + logger.info( + "WARMUP_COMPLETED path=%s round=%d/%d seconds=%.3f", + path, + index + 1, + args.warmup, + elapsed, + ) + + inference_times: list[float] = [] + result = None + for index in range(args.rounds): + provider.synchronize() + started = time.perf_counter() + result = pipeline.predict(str(path), **predict_kwargs) + provider.synchronize() + elapsed = time.perf_counter() - started + inference_times.append(elapsed) + logger.info( + "INFERENCE_COMPLETED path=%s round=%d/%d seconds=%.3f", + path, + index + 1, + args.rounds, + elapsed, + ) + + if not result: + raise RuntimeError("OCR pipeline 未返回图片结果") + + summary = _result_summary(result) + processing_seconds = time.perf_counter() - file_started + benchmark = { + "timestamp": datetime.now().astimezone().isoformat(), + **provider.metadata(), + "image_path": str(path), + "image": summary, + "warmup_seconds": [round(value, 3) for value in warmup_times], + "inference_seconds": { + "all": [round(value, 3) for value in inference_times], + "min": round(min(inference_times), 3), + "max": round(max(inference_times), 3), + "mean": round(statistics.fmean(inference_times), 3), + "median": round(statistics.median(inference_times), 3), + "stdev": round(statistics.pstdev(inference_times), 3), + }, + "gpu_memory": provider.gpu_memory(), + "processing_seconds": round(processing_seconds, 3), + "export_seconds": 0.0, + "file_total_seconds": 0.0, + } + output_dir = image_output_directory( + args.output, + path, + batch_root=batch_root, + recursive=args.recursive, + ) + export_started = time.perf_counter() + output_paths = save_image_ocr_outputs( + result[0], + output_dir, + input_path=path, + benchmark=benchmark, + ) + export_seconds = time.perf_counter() - export_started + total_seconds = time.perf_counter() - file_started + benchmark["export_seconds"] = round(export_seconds, 3) + benchmark["file_total_seconds"] = round(total_seconds, 3) + atomic_write_json(Path(output_paths["benchmark"]), benchmark) + if batch_root is None and args.benchmark_json: + explicit_benchmark = args.benchmark_json.expanduser().resolve() + atomic_write_json(explicit_benchmark, benchmark) + output_paths["explicit_benchmark"] = str(explicit_benchmark) + + logger.info( + "IMAGE_COMPLETED path=%s width=%s height=%s layout_boxes=%d parsed_blocks=%d inference_mean_seconds=%.3f export_seconds=%.3f file_total_seconds=%.3f output=%s benchmark=%s", + path, + summary["width"], + summary["height"], + summary["layout_boxes"], + summary["parsed_blocks"], + statistics.fmean(inference_times), + export_seconds, + total_seconds, + output_paths["output_dir"], + output_paths["benchmark"], + ) + if not args.no_result: + for index, block in enumerate(result[0]["parsing_res_list"], 1): + logger.info( + "OCR_BLOCK path=%s index=%d label=%s bbox=%s content=%s", + path, + index, + block.label, + block.bbox, + block.content.replace("\r", "").replace("\n", "\\n"), + ) + + return FileProcessResult( + path=path, + kind="image", + status="completed", + seconds=total_seconds, + details={**summary, **output_paths}, + ) + + +def process_pdf_file( + path: Path, + *, + args, + provider: PipelineProvider, + logger: logging.Logger, + batch_root: Path | None, +) -> FileProcessResult: + file_started = time.perf_counter() + logger.info("FILE_ROUTED path=%s kind=pdf mode=%s", path, args.pdf_mode) + output_root = pdf_output_root( + args.output, + path, + batch_root=batch_root, + recursive=args.recursive, + ) + manifest_exists = (output_root / safe_stem(path.stem) / "manifest.json").is_file() + # Directory jobs auto-resume existing PDF manifests so rerunning a batch is + # safe. Single-file jobs still require an explicit --resume. + resume = args.resume if batch_root is None else manifest_exists and not args.overwrite + preflight = preflight_pdf( + pdf_path=path, + output_root=output_root, + pages=args.pages, + dpi=args.dpi, + password=args.password, + resume=resume, + overwrite=args.overwrite, + ) + logger.info( + "PDF_PREFLIGHT_COMPLETED path=%s page_count=%d selected_pages=%d output=%s mode=%s", + path, + preflight["page_count"], + len(preflight["selected_pages"]), + preflight["document_dir"], + args.pdf_mode, + ) + policy = TextLayerPolicy( + min_chars=args.text_min_chars, + min_printable_ratio=args.text_min_printable_ratio, + min_content_ratio=args.text_min_content_ratio, + max_replacement_ratio=args.text_max_replacement_ratio, + min_chars_per_megapixel=args.text_min_density, + ) + if args.pdf_mode == "ocr": + provider.prepare() + summary = process_pdf( + provider=provider, + pdf_path=path, + output_root=output_root, + mode=args.pdf_mode, + text_policy=policy, + pages=args.pages, + dpi=args.dpi, + password=args.password, + resume=resume, + overwrite=args.overwrite, + keep_rendered=args.keep_rendered, + fail_fast=args.fail_fast, + 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 + }, + logger=logger, + ) + total_seconds = time.perf_counter() - file_started + logger.info( + "PDF_COMPLETED path=%s status=%s text_pages=%d ocr_pages=%d failed_pages=%s model_used=%s model_initialized_during_task=%s resume=%s file_total_seconds=%.3f output=%s", + path, + summary["status"], + summary["text_pages"], + summary["ocr_pages"], + summary["failed_pages"], + summary["model_used"], + summary["model_initialized_during_task"], + resume, + total_seconds, + summary["document_dir"], + ) + return FileProcessResult( + path=path, + kind="pdf", + status=summary["status"], + seconds=total_seconds, + details=summary, + exit_code=0 if not summary["failed_pages"] else 3, + ) + + +def process_single_file( + path: Path, + *, + args, + provider: PipelineProvider, + logger: logging.Logger, + project_root: Path, + run_image_warmup: bool, + batch_root: Path | None = None, +) -> FileProcessResult: + path = path.expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"文件不存在: {path}") + kind = detect_input_kind(path) + if kind == "image": + return process_image_file( + path, + args=args, + provider=provider, + logger=logger, + project_root=project_root, + run_warmup=run_image_warmup, + batch_root=batch_root, + ) + if kind == "pdf": + return process_pdf_file( + path, + args=args, + provider=provider, + logger=logger, + batch_root=batch_root, + ) + supported = ", ".join(sorted(SUPPORTED_EXTENSIONS)) + raise ValueError(f"不支持的文件类型: {path.suffix or '<无后缀>'}；支持: {supported}") + + +def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int: + program_started = time.perf_counter() + input_path = args.input.expanduser().resolve() + kind = detect_input_kind(input_path) + + if kind != "directory": + try: + result = process_single_file( + input_path, + args=args, + provider=provider, + logger=logger, + project_root=project_root, + run_image_warmup=True, + ) + except KeyboardInterrupt: + logger.warning("PROGRAM_INTERRUPTED input=%s resume_hint=--resume", input_path) + return 130 + except Exception as exc: + logger.exception("FILE_FAILED path=%s error=%s", input_path, exc) + return 1 + logger.info( + "PROGRAM_COMPLETED input=%s kind=%s status=%s file_seconds=%.3f program_total_seconds=%.3f", + result.path, + result.kind, + result.status, + result.seconds, + time.perf_counter() - program_started, + ) + return result.exit_code + + files = discover_supported_files( + input_path, + recursive=args.recursive, + output_root=args.output, + ) + if not files: + logger.error("NO_SUPPORTED_FILES directory=%s recursive=%s", input_path, args.recursive) + return 1 + logger.info( + "DIRECTORY_PLAN directory=%s recursive=%s files=%d image_files=%d pdf_files=%d", + input_path, + args.recursive, + len(files), + sum(path.suffix.lower() in IMAGE_EXTENSIONS for path in files), + sum(path.suffix.lower() in PDF_EXTENSIONS for path in files), + ) + + results: list[FileProcessResult] = [] + failures: list[dict[str, str]] = [] + image_warmup_pending = True + for index, path in enumerate(files, 1): + logger.info("DIRECTORY_PROGRESS_START progress=%d/%d path=%s", index, len(files), path) + try: + result = process_single_file( + path, + args=args, + provider=provider, + logger=logger, + project_root=project_root, + run_image_warmup=image_warmup_pending, + batch_root=input_path, + ) + results.append(result) + if result.kind == "image": + image_warmup_pending = False + if result.exit_code: + failures.append({"path": str(path), "error": result.status}) + except KeyboardInterrupt: + logger.warning("PROGRAM_INTERRUPTED path=%s progress=%d/%d", path, index, len(files)) + return 130 + except Exception as exc: + failures.append({"path": str(path), "error": f"{type(exc).__name__}: {exc}"}) + logger.exception("FILE_FAILED path=%s progress=%d/%d", path, index, len(files)) + if args.fail_fast: + break + logger.info("DIRECTORY_PROGRESS_END progress=%d/%d path=%s", index, len(files), path) + + image_results = [result for result in results if result.kind == "image"] + pdf_results = [result for result in results if result.kind == "pdf"] + total_file_seconds = sum(result.seconds for result in results) + program_total = time.perf_counter() - program_started + batch_manifest_path = ( + args.output.expanduser().resolve() + / "batches" + / f"{safe_stem(input_path.name or 'batch')}-{datetime.now():%Y%m%d-%H%M%S-%f}.json" + ) + batch_manifest = { + "input_directory": str(input_path), + "recursive": args.recursive, + "device": provider.resolved_device, + "discovered_files": len(files), + "completed_files": len(results), + "failed_files": len(failures), + "image_files": len(image_results), + "pdf_files": len(pdf_results), + "model_init_seconds": round(provider.model_init_seconds, 3), + "total_file_seconds": round(total_file_seconds, 3), + "program_total_seconds": round(program_total, 3), + "results": [ + { + "path": str(result.path), + "kind": result.kind, + "status": result.status, + "seconds": round(result.seconds, 3), + "exit_code": result.exit_code, + "outputs": result.details, + } + for result in results + ], + "failures": failures, + } + atomic_write_json(batch_manifest_path, batch_manifest) + logger.info( + "DIRECTORY_SUMMARY discovered=%d completed=%d failed=%d images_completed=%d pdfs_completed=%d total_file_seconds=%.3f program_total_seconds=%.3f model_init_seconds=%.3f manifest=%s", + len(files), + len(results), + len(failures), + len(image_results), + len(pdf_results), + total_file_seconds, + program_total, + provider.model_init_seconds, + batch_manifest_path, + ) + return 0 if not failures else 3 + + +def run_verify(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int: + started = time.perf_counter() + try: + provider.prepare() + paddle = provider._paddle + if provider.config.device == "gpu": + tensor = paddle.ones([1024, 1024], dtype="float32") + result = paddle.matmul(tensor, tensor) + provider.synchronize() + logger.info("GPU_SMOKE_TEST shape=%s", list(result.shape)) + else: + from paddle import core + + logger.info( + "CPU_SMOKE_TEST onednn=%s mkldnn=%s", + core.is_compiled_with_onednn(), + core.is_compiled_with_mkldnn(), + ) + logger.info( + "VERIFY_COMPLETED metadata=%s seconds=%.3f", + provider.metadata(), + time.perf_counter() - started, + ) + return 0 + except Exception: + logger.exception("VERIFY_FAILED") + return 1 diff --git a/ocr_logging.py b/ocr_app/logging_utils.py similarity index 100% rename from ocr_logging.py rename to ocr_app/logging_utils.py diff --git a/ocr_app/output.py b/ocr_app/output.py new file mode 100644 index 0000000..d4cc402 --- /dev/null +++ b/ocr_app/output.py @@ -0,0 +1,139 @@ +"""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 diff --git a/ocr_app/pdf.py b/ocr_app/pdf.py new file mode 100644 index 0000000..9e342f0 --- /dev/null +++ b/ocr_app/pdf.py @@ -0,0 +1,541 @@ +"""Hybrid PDF processing: extract usable text layers and OCR only when needed.""" + +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, Iterable + +import pypdfium2 as pdfium +from PIL import Image + +from .pdf_text import TextLayerPolicy, extract_page_text + +MANIFEST_VERSION = 2 +PAGE_SPEC_PATTERN = re.compile(r"^(\d+)(?:-(\d*)?)?$") +PDF_MODES = {"hybrid", "text", "ocr"} + + +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]: + 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 _save_markdown_image(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"无法保存 Markdown 图片: {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 _ocr_markdown(result: Any, document_dir: Path, page_number: int) -> str: + data = result.markdown + if "res" in data and isinstance(data["res"], dict): + data = data["res"] + text = str(data.get("markdown_texts", "")) + asset_dir = document_dir / "assets" / f"page-{page_number:04d}" + if asset_dir.exists(): + shutil.rmtree(asset_dir) + for index, (original_path, image_data) in enumerate((data.get("markdown_images") or {}).items(), 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)}{Path(source_name).suffix or '.png'}" + target = _save_markdown_image(image_data, target) + relative = Path(os.path.relpath(target, document_dir / "pages")).as_posix() + text = text.replace(original, relative).replace(str(original_path), relative) + return text.strip() + + +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(): + text = markdown_path.read_text(encoding="utf-8").replace("../assets/", "assets/") + source = record.get("source_type", "unknown") + markdown_parts.append(f"\n\n---\n\n## Page {page_number} ({source})\n\n{text.strip()}") + 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, + 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: + raise ValueError("旧版 manifest 不兼容混合模式，请使用 --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") + 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, + "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, + "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" + opened = time.perf_counter() + document = pdfium.PdfDocument(str(pdf_path), password=password) + pdf_open_seconds = time.perf_counter() - opened + logger.info("PDF_OPENED path=%s mode=%s seconds=%.3f dpi=%d", pdf_path, mode, pdf_open_seconds, dpi) + + try: + page_count = len(document) + selected_indexes = parse_page_spec(pages, page_count) + prepared = time.perf_counter() + 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, + run_metadata={"device": provider.resolved_device}, + ) + manifest_prepare_seconds = time.perf_counter() - prepared + 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)] + logger.info( + "TASK_PLAN mode=%s total_pages=%d selected_pages=%d completed_before=%d pending_pages=%d", + mode, page_count, len(selected_indexes), completed_before, len(pending), + ) + + current_run_times: list[float] = [] + 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 = state_save_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" + logger.info("PAGE_START page=%d position=%d/%d mode=%s", page_number, position, len(pending), mode) + 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" + logger.info( + "PAGE_ROUTED page=%d source=%s reason=%s text_chars=%d printable_ratio=%.4f content_ratio=%.4f density=%.3f text_extract_seconds=%.3f", + page_number, source_type, assessment.reason, assessment.non_whitespace_chars, + assessment.printable_ratio, assessment.content_ratio, assessment.chars_per_megapixel, + text_extract_seconds, + ) + + markdown_path, json_path = _page_paths(document_dir, page_number) + if source_type == "text": + markdown_text = extracted_text + payload = { + "res": { + "input_path": str(pdf_path), + "page_index": page_index, + "page_number": page_number, + "page_count": page_count, + "source_type": "text", + "text": extracted_text, + "text_layer": assessment_dict, + "width_points": width_points, + "height_points": height_points, + } + } + layout_boxes = 0 + parsed_blocks = 1 if extracted_text else 0 + 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("OCR pipeline 未返回结果") + result = results[0] + markdown_text = _ocr_markdown(result, document_dir, page_number) + payload = result.json + result_payload = payload.get("res", payload) + result_payload.update( + { + "input_path": str(pdf_path), + "page_index": page_index, + "page_number": page_number, + "page_count": page_count, + "source_type": "ocr", + "ocr_reason": assessment.reason if mode == "hybrid" else "forced_ocr_mode", + "text_layer": assessment_dict, + "render_dpi": dpi, + } + ) + layout_boxes = len(result.get("layout_det_res", {}).get("boxes", [])) + parsed_blocks = len(result.get("parsing_res_list", [])) + export_started = time.perf_counter() + atomic_write_text(markdown_path, markdown_text.rstrip() + "\n") + 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), + "layout_boxes": layout_boxes, + "parsed_blocks": parsed_blocks, + "completed_at": now_iso(), + } + current_run_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) + logger.warning("TASK_INTERRUPTED page=%d", page_number) + raise + except Exception as exc: + total_seconds = time.perf_counter() - started + manifest["pages"][str(page_number)] = { + "status": "failed", + "page_number": page_number, + "source_type": source_type, + "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), + "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() + + saved = time.perf_counter() + manifest["run_metadata"] = provider.metadata() + manifest["updated_at"] = now_iso() + atomic_write_json(manifest_path, manifest) + rebuild_combined_outputs(document_dir, manifest) + state_save_seconds = time.perf_counter() - saved + manifest["pages"][str(page_number)]["state_save_seconds"] = round(state_save_seconds, 3) + atomic_write_json(manifest_path, manifest) + average = sum(current_run_times) / len(current_run_times) if current_run_times else None + eta = average * (len(pending) - position) if average is not None else None + record = manifest["pages"][str(page_number)] + logger.info( + "PAGE_FINISHED page=%d status=%s source=%s text_extract_seconds=%.3f render_seconds=%.3f ocr_seconds=%.3f export_seconds=%.3f state_save_seconds=%.3f total_seconds=%.3f eta_seconds=%s progress=%d/%d", + page_number, record["status"], record.get("source_type"), text_extract_seconds, + render_seconds, ocr_seconds, export_seconds, state_save_seconds, + record.get("total_seconds", 0.0), f"{eta:.3f}" if eta is not None else "unknown", + 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) + timing_keys = ("text_extract_seconds", "render_seconds", "ocr_seconds", "export_seconds", "state_save_seconds", "total_seconds") + totals = {key: sum(record.get(key, 0.0) for record in records) for key in timing_keys} + finalize_started = time.perf_counter() + 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 + ), + "model_available": provider.model_init_seconds > 0, + "timing": { + "pdf_open_seconds": round(pdf_open_seconds, 3), + "manifest_prepare_seconds": round(manifest_prepare_seconds, 3), + "text_extract_total_seconds": round(totals["text_extract_seconds"], 3), + "render_total_seconds": round(totals["render_seconds"], 3), + "ocr_total_seconds": round(totals["ocr_seconds"], 3), + "export_total_seconds": round(totals["export_seconds"], 3), + "state_save_total_seconds": round(totals["state_save_seconds"], 3), + "page_total_seconds": round(totals["total_seconds"], 3), + "model_init_seconds": round(provider.model_init_seconds, 3), + "finalize_seconds": 0.0, + "task_total_seconds": 0.0, + }, + } + 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 mode=%s selected_pages=%d text_pages=%d ocr_pages=%d failed_pages=%s model_used=%s model_initialized_during_task=%s model_available=%s model_init_seconds=%.3f text_extract_total_seconds=%.3f render_total_seconds=%.3f ocr_total_seconds=%.3f task_total_seconds=%.3f", + manifest["status"], mode, len(selected_indexes), text_pages, ocr_pages, failed_pages, + manifest["summary"]["model_used"], + manifest["summary"]["model_initialized_during_task"], + manifest["summary"]["model_available"], + provider.model_init_seconds, + totals["text_extract_seconds"], totals["render_seconds"], totals["ocr_seconds"], task_total, + ) + return {"document_dir": str(document_dir), "manifest_path": str(manifest_path), "status": manifest["status"], **manifest["summary"]} + finally: + document.close() diff --git a/ocr_app/pdf_text.py b/ocr_app/pdf_text.py new file mode 100644 index 0000000..0046eef --- /dev/null +++ b/ocr_app/pdf_text.py @@ -0,0 +1,127 @@ +"""PDF text-layer extraction and quality assessment for hybrid OCR.""" + +from __future__ import annotations + +import re +import unicodedata +from dataclasses import asdict, dataclass +from typing import Any + + +@dataclass +class TextLayerPolicy: + min_chars: int = 50 + min_printable_ratio: float = 0.85 + min_content_ratio: float = 0.60 + max_replacement_ratio: float = 0.02 + min_chars_per_megapixel: float = 25.0 + + +@dataclass +class TextLayerAssessment: + usable: bool + reason: str + raw_chars: int + non_whitespace_chars: int + printable_ratio: float + content_ratio: float + replacement_ratio: float + chars_per_megapixel: float + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def normalize_text(text: str) -> str: + text = text.replace("\x00", "") + text = text.replace("\r\n", "\n").replace("\r", "\n") + lines = [re.sub(r"[ \t]+", " ", line).strip() for line in text.splitlines()] + compact_lines: list[str] = [] + previous_blank = False + for line in lines: + if line: + compact_lines.append(line) + previous_blank = False + elif compact_lines and not previous_blank: + compact_lines.append("") + previous_blank = True + return "\n".join(compact_lines).strip() + + +def _is_content_character(character: str) -> bool: + if character.isalnum(): + return True + code = ord(character) + return ( + 0x3400 <= code <= 0x4DBF + or 0x4E00 <= code <= 0x9FFF + or 0xF900 <= code <= 0xFAFF + or 0x3040 <= code <= 0x30FF + or 0xAC00 <= code <= 0xD7AF + ) + + +def assess_text_layer( + text: str, + *, + width_points: float, + height_points: float, + policy: TextLayerPolicy, +) -> TextLayerAssessment: + normalized = normalize_text(text) + compact = [character for character in normalized if not character.isspace()] + count = len(compact) + if count == 0: + return TextLayerAssessment(False, "empty_text_layer", len(text), 0, 0.0, 0.0, 0.0, 0.0) + + printable = sum(character.isprintable() and unicodedata.category(character) != "Cc" for character in compact) + content = sum(_is_content_character(character) for character in compact) + replacements = sum(character in {"\ufffd", "�"} for character in compact) + page_megapixels = max((width_points * height_points) / 1_000_000.0, 0.01) + printable_ratio = printable / count + content_ratio = content / count + replacement_ratio = replacements / count + density = count / page_megapixels + + checks = ( + (count >= policy.min_chars, "too_few_characters"), + (printable_ratio >= policy.min_printable_ratio, "low_printable_ratio"), + (content_ratio >= policy.min_content_ratio, "low_content_ratio"), + (replacement_ratio <= policy.max_replacement_ratio, "high_replacement_ratio"), + (density >= policy.min_chars_per_megapixel, "low_text_density"), + ) + reason = "usable_text_layer" + usable = True + for passed, failure_reason in checks: + if not passed: + usable = False + reason = failure_reason + break + + return TextLayerAssessment( + usable=usable, + reason=reason, + raw_chars=len(text), + non_whitespace_chars=count, + printable_ratio=round(printable_ratio, 4), + content_ratio=round(content_ratio, 4), + replacement_ratio=round(replacement_ratio, 4), + chars_per_megapixel=round(density, 3), + ) + + +def extract_page_text(page: Any, policy: TextLayerPolicy) -> tuple[str, TextLayerAssessment]: + text_page = page.get_textpage() + try: + raw_text = text_page.get_text_bounded() + finally: + text_page.close() + width, height = page.get_size() + normalized = normalize_text(raw_text) + assessment = assess_text_layer( + normalized, + width_points=width, + height_points=height, + policy=policy, + ) + return normalized, assessment diff --git a/ocr_app/runtime.py b/ocr_app/runtime.py new file mode 100644 index 0000000..4131a64 --- /dev/null +++ b/ocr_app/runtime.py @@ -0,0 +1,170 @@ +"""Device validation and lazy PaddleOCR-VL pipeline creation.""" + +from __future__ import annotations + +import logging +import os +import platform +import time +from dataclasses import dataclass +from typing import Any + + +@dataclass +class RuntimeConfig: + device: str + threads: int | None = None + device_id: int = 0 + + +class PipelineProvider: + """Create the large OCR pipeline only when a command actually needs it.""" + + def __init__(self, config: RuntimeConfig, logger: logging.Logger): + self.config = config + self.logger = logger + self._pipeline: Any | None = None + self._paddle: Any | None = None + self._device_name: str | None = None + self.import_seconds = 0.0 + self.setup_seconds = 0.0 + self.model_init_seconds = 0.0 + + @property + def resolved_device(self) -> str: + return "cpu" if self.config.device == "cpu" else f"gpu:{self.config.device_id}" + + def prepare(self) -> None: + """Validate and configure Paddle without loading the OCR model.""" + if self._paddle is not None: + return + + started = time.perf_counter() + try: + import paddle + except ImportError as exc: + package = "paddlepaddle" if self.config.device == "cpu" else "paddlepaddle-gpu" + raise RuntimeError(f"当前子项目未安装 {package}") from exc + self.import_seconds = time.perf_counter() - started + self._paddle = paddle + + setup_started = time.perf_counter() + if self.config.device == "cpu": + from paddle import core + + total_cores = os.cpu_count() or 4 + threads = self.config.threads or max(1, total_cores - 2) + if threads < 1: + raise ValueError("CPU 线程数必须大于等于 1") + self.config.threads = threads + core.set_num_threads(threads) + self._device_name = platform.processor() or "CPU" + paddle.set_device("cpu") + self.logger.info( + "CPU_CONFIGURED threads=%d total_cores=%d reserved_cores=%d", + threads, + total_cores, + max(0, total_cores - threads), + ) + else: + 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 self.config.device_id < 0 or self.config.device_id >= device_count: + raise RuntimeError( + f"GPU {self.config.device_id} 不存在，当前检测到 {device_count} 个设备" + ) + paddle.set_device(self.resolved_device) + paddle.device.cuda.synchronize(self.config.device_id) + try: + self._device_name = paddle.device.cuda.get_device_name(self.config.device_id) + except Exception: + self._device_name = "unknown" + self.logger.info( + "GPU_CONFIGURED device=%s device_name=%s device_count=%d", + self.resolved_device, + self._device_name, + device_count, + ) + self.setup_seconds = time.perf_counter() - setup_started + self.logger.info( + "RUNTIME_PREPARED device=%s paddle_version=%s import_seconds=%.3f setup_seconds=%.3f", + self.resolved_device, + paddle.__version__, + self.import_seconds, + self.setup_seconds, + ) + + def get(self): + self.prepare() + if self._pipeline is None: + self.logger.info( + "MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=%s", + self.resolved_device, + ) + started = time.perf_counter() + from paddleocr import PaddleOCRVL + + self._pipeline = PaddleOCRVL( + pipeline_version="v1.6", + device=self.resolved_device, + ) + self.synchronize() + self.model_init_seconds = time.perf_counter() - started + self.logger.info( + "MODEL_INITIALIZED seconds=%.3f device=%s", + self.model_init_seconds, + self.resolved_device, + ) + return self._pipeline + + def synchronize(self) -> None: + if self.config.device == "gpu" and self._paddle is not None: + self._paddle.device.cuda.synchronize(self.config.device_id) + + def gpu_memory(self) -> dict[str, float | None]: + stats: dict[str, float | None] = { + "allocated_mb": None, + "reserved_mb": None, + "max_allocated_mb": None, + "max_reserved_mb": None, + } + if self.config.device != "gpu" or self._paddle is None: + return stats + functions = { + "allocated_mb": "memory_allocated", + "reserved_mb": "memory_reserved", + "max_allocated_mb": "max_memory_allocated", + "max_reserved_mb": "max_memory_reserved", + } + for key, name in functions.items(): + function = getattr(self._paddle.device.cuda, name, None) + if function is None: + continue + try: + stats[key] = round( + float(function(self.config.device_id)) / (1024**2), 2 + ) + except Exception: + pass + return stats + + def metadata(self) -> dict[str, Any]: + paddle_version = self._paddle.__version__ if self._paddle is not None else None + return { + "device": self.resolved_device, + "device_name": self._device_name, + "cpu_threads": self.config.threads if self.config.device == "cpu" else None, + "python_version": platform.python_version(), + "platform": platform.platform(), + "paddle_version": paddle_version, + "pipeline_version": "v1.6", + "runtime_import_seconds": round(self.import_seconds, 3), + "runtime_setup_seconds": round(self.setup_seconds, 3), + "model_init_seconds": round(self.model_init_seconds, 3), + } diff --git a/pdf_ocr.py b/pdf_ocr.py deleted file mode 100644 index b5df4d5..0000000 --- a/pdf_ocr.py +++ /dev/null @@ -1,182 +0,0 @@ -"""CPU entry point for page-by-page PaddleOCR-VL PDF recognition.""" - -from __future__ import annotations - -import argparse -import os -import platform -import time -from pathlib import Path - -from ocr_logging import default_log_path, setup_run_logger -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 最大输入像素参数") - parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") - parser.add_argument("--verbose", action="store_true", help="输出详细日志") - return parser.parse_args() - - -def main() -> int: - program_started = time.perf_counter() - args = parse_args() - log_file = args.log_file or default_log_path( - PROJECT_ROOT, - "pdf", - args.pdf.stem, - device="cpu", - ) - logger = setup_run_logger("ocr.pdf.cpu", log_file, verbose=args.verbose) - logger.info( - "PROGRAM_STARTED input=%s output=%s pages=%s dpi=%d resume=%s overwrite=%s keep_rendered=%s fail_fast=%s", - args.pdf, - args.output, - args.pages or "all", - args.dpi, - args.resume, - args.overwrite, - args.keep_rendered, - args.fail_fast, - ) - 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: - logger.error("INVALID_ARGUMENT threads=%d", threads) - return 2 - - try: - preflight_started = time.perf_counter() - 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, - ) - preflight_seconds = time.perf_counter() - preflight_started - except Exception as exc: - logger.error("PREFLIGHT_FAILED type=%s error=%s", type(exc).__name__, exc, exc_info=args.verbose) - return 1 - - logger.info( - "PREFLIGHT_COMPLETED seconds=%.3f page_count=%d selected_pages=%d document_dir=%s", - preflight_seconds, - preflight["page_count"], - len(preflight["selected_pages"]), - preflight["document_dir"], - ) - - import_started = time.perf_counter() - from paddle import core - from paddleocr import PaddleOCRVL - import_seconds = time.perf_counter() - import_started - - core.set_num_threads(threads) - logger.info( - "RUNTIME_READY import_seconds=%.3f threads=%d total_cores=%d reserved_cores=%d", - import_seconds, - threads, - total_cores, - max(0, total_cores - threads), - ) - logger.info("MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=cpu") - init_started = time.perf_counter() - pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu") - init_seconds = time.perf_counter() - init_started - logger.info("MODEL_INITIALIZED seconds=%.3f pipeline_version=v1.6 device=cpu", init_seconds) - - 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", - "preflight_seconds": round(preflight_seconds, 3), - "runtime_import_seconds": round(import_seconds, 3), - "log_file": str(log_file.resolve()), - } - - 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, - logger=logger, - ) - except KeyboardInterrupt: - logger.warning( - "PROGRAM_INTERRUPTED total_seconds=%.3f resume_hint=--resume", - time.perf_counter() - program_started, - ) - return 130 - except Exception as exc: - logger.exception( - "PROGRAM_FAILED type=%s error=%s total_seconds=%.3f", - type(exc).__name__, - exc, - time.perf_counter() - program_started, - ) - return 1 - - program_total = time.perf_counter() - program_started - timing = summary.get("timing", {}) - logger.info( - "PROGRAM_COMPLETED status=%s completed_pages=%d selected_pages=%d failed_pages=%s model_init_seconds=%.3f pdf_task_seconds=%.3f program_total_seconds=%.3f output=%s log=%s", - summary["status"], - summary["completed_pages"], - summary["selected_pages"], - summary["failed_pages"], - init_seconds, - timing.get("task_total_seconds", 0.0), - program_total, - summary["document_dir"], - log_file.resolve(), - ) - 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 deleted file mode 100644 index f924c52..0000000 --- a/pdf_ocr_core.py +++ /dev/null @@ -1,658 +0,0 @@ -"""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() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7d6fded --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,6 @@ +from pathlib import Path +import sys + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) diff --git a/tests/test_launcher.py b/tests/test_launcher.py new file mode 100644 index 0000000..7a939b2 --- /dev/null +++ b/tests/test_launcher.py @@ -0,0 +1,13 @@ +from ocr import _requested_device + + +def test_default_device(): + assert _requested_device(["verify"]) == "cpu" + + +def test_requested_gpu(): + assert _requested_device(["verify", "--device", "gpu"]) == "gpu" + + +def test_requested_gpu_equals_syntax(): + assert _requested_device(["verify", "--device=gpu"]) == "gpu" diff --git a/tests/test_output_routing.py b/tests/test_output_routing.py new file mode 100644 index 0000000..6dc3578 --- /dev/null +++ b/tests/test_output_routing.py @@ -0,0 +1,126 @@ +import json +from argparse import Namespace +from pathlib import Path + +from PIL import Image + +from ocr_app.commands import process_image_file +from ocr_app.logging_utils import setup_run_logger +from ocr_app.output import image_output_directory, pdf_output_root + + +class FakeBlock: + label = "text" + bbox = [0, 0, 10, 10] + content = "hello OCR" + image = None + + +class FakeResult(dict): + @property + def markdown(self): + return {"markdown_texts": "hello OCR", "markdown_images": {}} + + @property + def json(self): + return {"res": {"input_path": self["input_path"]}} + + +class FakePipeline: + def predict(self, path): + return [ + FakeResult( + input_path=path, + width=20, + height=10, + layout_det_res={"boxes": [{}]}, + parsing_res_list=[FakeBlock()], + ) + ] + + +class FakeProvider: + class Config: + device = "cpu" + + config = Config() + model_init_seconds = 0.01 + + def get(self): + return FakePipeline() + + def synchronize(self): + pass + + def metadata(self): + return {"device": "cpu", "model_init_seconds": self.model_init_seconds} + + def gpu_memory(self): + return {} + + +def make_args(output): + return Namespace( + warmup=0, + rounds=1, + output=output, + recursive=False, + benchmark_json=None, + no_result=True, + ) + + +def test_single_image_generates_output_files(tmp_path): + image = tmp_path / "card.jpg" + Image.new("RGB", (20, 10), "white").save(image) + logger = setup_run_logger("test.image.output", tmp_path / "run.log", console=False) + + result = process_image_file( + image, + args=make_args(tmp_path / "outputs"), + provider=FakeProvider(), + logger=logger, + project_root=tmp_path, + run_warmup=True, + batch_root=None, + ) + + output_dir = Path(result.details["output_dir"]) + assert output_dir == tmp_path / "outputs" / "images" / "card_jpg" + assert (output_dir / "result.md").read_text("utf-8").strip() == "hello OCR" + assert (output_dir / "result.txt").read_text("utf-8").strip() == "hello OCR" + data = json.loads((output_dir / "result.json").read_text("utf-8")) + assert data["res"]["source_type"] == "image_ocr" + benchmark = json.loads((output_dir / "benchmark.json").read_text("utf-8")) + assert benchmark["file_total_seconds"] >= benchmark["processing_seconds"] + + +def test_recursive_output_paths_preserve_relative_directories(tmp_path): + batch_root = tmp_path / "input" + image = batch_root / "sub" / "same.png" + pdf = batch_root / "other" / "same.pdf" + image.parent.mkdir(parents=True) + pdf.parent.mkdir(parents=True) + output = tmp_path / "outputs" + + assert image_output_directory( + output, + image, + batch_root=batch_root, + recursive=True, + ) == output / "images" / "sub" / "same_png" + assert pdf_output_root( + output, + pdf, + batch_root=batch_root, + recursive=True, + ) == output / "pdfs" / "other" + + +def test_image_extensions_do_not_collide(tmp_path): + output = tmp_path / "outputs" + png = tmp_path / "same.png" + jpg = tmp_path / "same.jpg" + assert image_output_directory(output, png, batch_root=None, recursive=False) != image_output_directory( + output, jpg, batch_root=None, recursive=False + ) diff --git a/tests/test_page_spec.py b/tests/test_page_spec.py new file mode 100644 index 0000000..8837200 --- /dev/null +++ b/tests/test_page_spec.py @@ -0,0 +1,13 @@ +import pytest + +from ocr_app.pdf import parse_page_spec + + +def test_page_ranges(): + assert parse_page_spec("1-2,4,6-", 7) == [0, 1, 3, 5, 6] + + +@pytest.mark.parametrize("value", ["0", "3-2", "1-a", "8"]) +def test_invalid_page_ranges(value): + with pytest.raises(ValueError): + parse_page_spec(value, 7) diff --git a/tests/test_pdf_hybrid.py b/tests/test_pdf_hybrid.py new file mode 100644 index 0000000..a1022b7 --- /dev/null +++ b/tests/test_pdf_hybrid.py @@ -0,0 +1,102 @@ +import json +from pathlib import Path + +from PIL import Image + +from ocr_app.pdf import process_pdf +from ocr_app.pdf_text import TextLayerPolicy + + +class FakeBlock: + label = "text" + bbox = [0, 0, 10, 10] + content = "mock OCR" + image = None + + +class FakeResult(dict): + @property + def markdown(self): + return {"markdown_texts": "mock OCR", "markdown_images": {}} + + @property + def json(self): + return {"res": {"input_path": self["input_path"]}} + + +class FakePipeline: + def predict(self, path, **kwargs): + return [ + FakeResult( + input_path=path, + width=144, + height=144, + layout_det_res={"boxes": [{}]}, + parsing_res_list=[FakeBlock()], + ) + ] + + +class FakeProvider: + resolved_device = "cpu" + + def __init__(self): + self.model_init_seconds = 0.0 + self.get_calls = 0 + + def get(self): + self.get_calls += 1 + self.model_init_seconds = 0.01 + return FakePipeline() + + def synchronize(self): + pass + + def metadata(self): + return { + "device": "cpu", + "model_init_seconds": self.model_init_seconds, + } + + +def test_electronic_pdf_does_not_load_model(tmp_path): + project_root = Path(__file__).resolve().parent.parent + source = next((project_root / "data" / "documents").glob("*.pdf")) + provider = FakeProvider() + + result = process_pdf( + provider=provider, + pdf_path=source, + output_root=tmp_path, + mode="hybrid", + pages="1", + text_policy=TextLayerPolicy(), + ) + + assert result["text_pages"] == 1 + assert result["ocr_pages"] == 0 + assert not result["model_used"] + assert not result["model_initialized_during_task"] + assert provider.get_calls == 0 + + +def test_scanned_pdf_falls_back_to_ocr(tmp_path): + source = tmp_path / "scan.pdf" + Image.new("RGB", (72, 72), "white").save(source) + provider = FakeProvider() + + result = process_pdf( + provider=provider, + pdf_path=source, + output_root=tmp_path / "output", + mode="hybrid", + text_policy=TextLayerPolicy(), + ) + + manifest = json.loads(Path(result["manifest_path"]).read_text(encoding="utf-8")) + assert result["text_pages"] == 0 + assert result["ocr_pages"] == 1 + assert result["model_used"] + assert result["model_initialized_during_task"] + assert provider.get_calls == 1 + assert manifest["pages"]["1"]["routing_reason"] == "empty_text_layer" diff --git a/tests/test_pdf_text.py b/tests/test_pdf_text.py new file mode 100644 index 0000000..a4a9105 --- /dev/null +++ b/tests/test_pdf_text.py @@ -0,0 +1,39 @@ +from ocr_app.pdf_text import TextLayerPolicy, assess_text_layer, normalize_text + + +def test_normalize_text(): + assert normalize_text("a b\r\n\r\n\r\nc") == "a b\n\nc" + + +def test_usable_text_layer(): + text = "有效电子文档内容 123 " * 20 + result = assess_text_layer( + text, + width_points=595, + height_points=842, + policy=TextLayerPolicy(), + ) + assert result.usable + assert result.reason == "usable_text_layer" + + +def test_empty_text_layer_routes_to_ocr(): + result = assess_text_layer( + "", + width_points=595, + height_points=842, + policy=TextLayerPolicy(), + ) + assert not result.usable + assert result.reason == "empty_text_layer" + + +def test_short_text_layer_routes_to_ocr(): + result = assess_text_layer( + "页码 1", + width_points=595, + height_points=842, + policy=TextLayerPolicy(min_chars=50), + ) + assert not result.usable + assert result.reason == "too_few_characters"