refactor:调整代码结构,将cpu模式收纳进子目录,统一入口,根据文件后缀和目录自动判断使用路由(pdf、图片和批量)

This commit is contained in:
kuuhaku 2026-07-16 16:20:55 +08:00
parent 406845930b
commit 2988521c41
44 changed files with 2522 additions and 2360 deletions

9
.gitignore vendored
View File

@ -8,15 +8,22 @@ wheels/
# Virtual environments # Virtual environments
.venv .venv
cpu/.venv/
gpu/.venv/
gpu/.gpu-ready
# Generated benchmark results # Generated benchmark results
benchmarks/cpu/*.json
benchmarks/gpu/*.json benchmarks/gpu/*.json
!benchmarks/cpu/.gitkeep
!benchmarks/gpu/.gitkeep !benchmarks/gpu/.gitkeep
# OCR outputs # OCR outputs
outputs/ outputs/
# Generated structured logs (legacy logs directly under logs/ remain tracked) # Generated structured logs (legacy logs directly under logs/ remain tracked)
logs/single/ logs/input/
logs/verify/
logs/image/
logs/batch/ logs/batch/
logs/pdf/ logs/pdf/

706
README.md
View File

@ -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/ ocr-VL1.6/
├── main.py # CPU 单图 OCR + Benchmark ├── ocr.py # 唯一用户入口,自动选择 CPU/GPU 子环境
├── batch_ocr.py # CPU 批量图片 OCR系统友好的多进程版本 ├── ocr_app/ # CPU/GPU 共享业务代码
├── pdf_ocr.py # CPU PDF OCR逐页、可恢复 │ ├── cli.py # image/batch/pdf/verify 命令
├── pdf_ocr_core.py # CPU/GPU 共用的 PDF 渲染、恢复和导出逻辑 │ ├── commands.py # 命令实现
├── ocr_logging.py # CPU/GPU 共用的 UTF-8 结构化日志工具 │ ├── runtime.py # 设备验证、模型延迟加载
├── pyproject.toml # CPU 项目依赖 │ ├── pdf.py # 混合 PDF、断点续传、导出
├── uv.lock # CPU 锁文件 │ ├── pdf_text.py # 文本层提取与质量评估
├── gpu/ # 独立 GPU 子项目 │ └── logging_utils.py # UTF-8 结构化日志
│ ├── main.py # GPU 单图 Benchmark ├── cpu/
│ ├── pdf_ocr.py # GPU PDF OCR复用公共核心 │ ├── pyproject.toml # CPU 独立依赖
│ ├── verify_env.py # CUDA 环境与计算验证 │ ├── uv.lock # CPU 独立锁文件
│ ├── setup_env.py # 按 CUDA Wheel 类型创建环境 │ ├── .python-version # Python 3.13
│ └── runner.py # 统一入口的 CPU 执行器
├── gpu/
│ ├── pyproject.toml # GPU 独立依赖 │ ├── pyproject.toml # GPU 独立依赖
│ ├── .python-version # GPU 使用 Python 3.11 │ ├── .python-version # Python 3.11
│ └── README.md # GPU 安装与运行说明 │ ├── setup_env.py # CUDA Wheel 安装脚本
├── benchmarks/ │ └── runner.py # 统一入口的 GPU 执行器
│ └── gpu/ # GPU Benchmark JSON 输出目录 ├── data/
├── images/ │ ├── images/ # 测试图片
│ └── 手写01.png # 测试图片手写中文1758×646 │ └── documents/ # 测试 PDF
└── README.md ├── outputs/ # PDF 输出git ignored
├── benchmarks/ # 单图 Benchmark JSON
├── logs/ # UTF-8 运行日志
└── tests/ # 共享逻辑测试
``` ```
## 技术栈 根目录不再保存 Paddle 虚拟环境或 Python 项目依赖。CPU 与 GPU 分别使用 `cpu/.venv``gpu/.venv`
| 组件 | 版本 | 说明 | ## 安装
| ---------------- | ----- | ---------------------------------------- |
| PythonCPU | 3.13 | 根目录独立环境 |
| PythonGPU | 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/` ### CPU
## 快速开始
### 前提条件
- Python ≥ 3.13
- [uv](https://github.com/astral-sh/uv) 包管理器
### 安装
```bash ```bash
uv sync uv sync --project cpu
``` ```
### 运行 ### GPU
根据目标机器 CUDA/驱动和 PaddlePaddle 官方兼容表选择 Wheel
```bash ```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 python gpu/setup_env.py --cuda cu118 --dry-run
# 在目标 NVIDIA GPU 机器创建 gpu/.venv根据官方兼容表选择 cu118 或 cu126
python gpu/setup_env.py --cuda cu118 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 ```bash
# 指定日志文件 python gpu/setup_env.py --cuda cu126
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
``` ```
日志文件使用 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 ```bash
# 处理整个 PDF默认 DPI 144 python ocr.py <文件或目录> --device cpu|gpu
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 线程数;长任务建议保留 12 个核心给系统
uv run python pdf_ocr.py documents/sample.pdf --threads 18
``` ```
### GPU 使用 程序自动判断输入类型:
GPU 入口与 CPU 入口使用相同的 PDF 核心逻辑,但必须在 `gpu/.venv` 和 NVIDIA CUDA GPU 上运行: | 输入 | 自动路由 |
|------|----------|
| `.png/.jpg/.jpeg/.bmp/.tif/.tiff/.webp` | 图片 OCR + Benchmark |
| `.pdf` | PDF 混合文本提取/OCR |
| 目录 | 发现其中的图片和 PDF逐个调用同一个单文件路由器 |
不支持的文件后缀会给出明确错误;目录模式会自动忽略不支持的文件。
### 验证环境
```bash ```bash
uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf \ python ocr.py verify --device cpu
--device-id 0 \ python ocr.py verify --device gpu
--pages "1-10" \
--dpi 144
``` ```
当前机器无 NVIDIA 独立显卡,因此 GPU PDF 入口仅完成静态检查,尚未实机验证。 ### 单张图片
### 页码语法 ```bash
python ocr.py data/images/手写01.png --device cpu
```
| 参数 | 含义 | 多轮 Benchmark
|------|------|
| `1` | 仅第 1 页 |
| `1-5` | 第 15 页 |
| `10-` | 第 10 页到最后一页 |
| `1-5,8,10-` | 多个页码范围组合 |
用户页码从 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 ```text
outputs/ outputs/
└── sample/ ├── images/
├── manifest.json # 任务配置、页状态、耗时和错误 │ └── <相对目录>/<图片名_扩展名>/
├── document.md # 合并后的 Markdown │ ├── result.md # Markdown 结果
├── document.json # 合并后的 JSON │ ├── result.txt # 纯文本结果
│ ├── result.json # PaddleOCR 结构化结果
│ ├── benchmark.json # 模型/推理/导出耗时
│ └── assets/ # 识别结果中的图片资源(如有)
├── pdfs/
│ └── <相对目录>/<PDF名>/
│ ├── manifest.json
│ ├── document.md
│ ├── document.json
│ └── pages/
└── batches/
└── <目录名>-<时间戳>.json # 批量任务汇总
```
例如:
```bash
python ocr.py data/images/名片01.jpg --device cpu
```
会生成:
```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` 相互覆盖。
重构前生成的 PDF 结果可能仍直接位于 `outputs/<PDF名>/`。这些旧结果不会自动删除;新任务统一写入 `outputs/pdfs/<PDF名>/`,确认不再需要后可手动迁移或清理。
目录任务再次运行时:
- 已存在的 PDF manifest 自动断点续传
- 新加入的 PDF 自动创建新任务
- 图片重新识别,并使用原子写入覆盖对应输出文件
- `--overwrite` 会强制 PDF 重新处理
- 每次目录任务都会生成新的 `outputs/batches/*.json` 汇总
旧命令前缀 `image`、`pdf`、`batch` 暂时兼容,例如 `python ocr.py image a.png`,但推荐直接传入路径。
## PDF 混合模式
### 默认hybrid
```bash
python ocr.py data/documents/sample.pdf --device cpu
```
每页流程:
```text
读取 PDF 文本层
文本质量评估
┌────┴────┐
│ │
有效文本 无效/不足
│ │
直接保存 渲染页面 → PaddleOCR-VL
└────┬────┘
逐页 Markdown/JSON + 合并文档
```
如果整份 PDF 都具有有效文本层,模型完全不会加载。实测项目中的电子合同 PDF单页混合处理约 `0.06s`,且 `model_used=false`
### 强制文本模式
```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
# 第 15 页、第 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/<PDF名>/
├── manifest.json
├── document.md
├── document.json
├── pages/ ├── pages/
│ ├── page-0001.md │ ├── page-0001.md
│ ├── page-0001.json │ ├── page-0001.json
│ └── ... │ └── ...
├── assets/ # 表格、图片等 Markdown 资源 ├── assets/
└── rendered/ # 仅使用 --keep-rendered 时保留 └── rendered/ # 仅 --keep-rendered 时存在
``` ```
默认不保留中间渲染 PNGOCR 完成后会删除临时图。每页 JSON 会将 `input_path` 恢复为原 PDF 路径,并记录 `page_index`、`page_number`、`page_count` 和 `render_dpi` 合并 Markdown 标题会标明页面来源:
### 恢复与错误处理 ```markdown
## Page 1 (text)
- 输出目录已存在时,必须显式使用 `--resume``--overwrite` ## Page 2 (ocr)
- `--resume` 会校验 PDF SHA-256 和 DPI防止接续到错误任务
- 单页失败默认写入 manifest 并继续后续页面
- `--fail-fast` 可在第一页失败后立即停止
- `Ctrl+C` 会保存当前 manifest下次使用 `--resume` 继续
- 逐页文件和 manifest 使用临时文件替换,降低中途退出造成文件损坏的概率
### DPI 建议
| 文档类型 | 建议 DPI |
|----------|---------:|
| 普通打印文字 | 120144 |
| 小字号文档 | 150200 |
| 手写或低质量扫描件 | 200250 |
CPU 当前单图实测约 162 秒。长 PDF 总时间可粗略按 `待处理页数 × 单页耗时` 估算,因此建议先用 `--pages "1"` 测试效果和耗时再扩大页码范围。DPI 越高通常越慢,不建议默认使用 300 DPI。
## 工作原理
`PaddleOCRVL` pipeline 分两阶段:
```
输入图片 → [PP-DocLayoutV3 版面检测] → [PaddleOCR-VL-1.6-0.9B 文字识别] → 结构化输出
``` ```
1. **版面检测** — PP-DocLayoutV3 检测页面中的文本块区域(坐标 + 标签 + 置信度) Manifest 汇总包括:
2. **OCR 识别** — PaddleOCR-VL-1.6-0.9BGQA 架构视觉语言模型)逐块识别文字
3. **结果输出** — 返回 `PaddleOCRVLResult`,包含布局信息和识别文本
### 输出结构 ```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 常用参数
| ---------------------- | ---------------------- | ------------------------------------------------------------ |
| `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 ```bash
# 2 进程并行(安全默认值) python ocr.py sample.pdf \
uv run python batch_ocr.py images/ --device cpu \
--pdf-mode hybrid \
# 4 进程并行(需 32GB+ RAM --pages "1-10" \
uv run python batch_ocr.py images/ --workers 4 --dpi 144 \
--threads 18 \
# 自定义错峰窗口(值越大内存峰值越低,但总耗时增加) --keep-rendered \
uv run python batch_ocr.py images/ --workers 4 --stagger 30 --log-file logs/pdf-sample.log
``` ```
| 配置 | 适用场景 | 理论加速比 | 内存开销 | 实际限制 | DPI 建议:
| -------------------- | -------- | --------------- | ---------- | -------------- |
| `set_num_threads(N)` | 单张图片 | ~1.5x | 无额外开销 | 自回归解码瓶颈 |
| `batch_ocr.py` | 批量多图 | ~NxN=进程数) | N × 2GB | 内存/内存带宽,需错峰避免打满系统 |
> ⚠️ 每个进程独立加载模型(~2GB32GB RAM 建议从 `--workers 2` 开始测试。默认值是相对保守配置,但是否稳定仍取决于可用内存、散热、后台应用和图片复杂度。 | 文档 | DPI |
|------|----:|
| 普通打印文字 | 120144 |
| 小字号 | 150200 |
| 手写/低质量扫描件 | 200250 |
--- DPI 只影响需要 OCR 的页面,不影响直接提取文本层的页面。
### 迭代 3独立 GPU 子项目(待实机验证) ## 日志
为避免 `paddlepaddle``paddlepaddle-gpu` 相互覆盖,在同一仓库新增 `gpu/` 子项目,使用独立 Python、虚拟环境、依赖配置和锁文件。 统一日志格式:
已完成: ```text
2026-07-16 15:07:45 | INFO | pid=33552 | PAGE_ROUTED page=1 source=text reason=usable_text_layer
```
- `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 和安全退出检查
尚未验证: ```text
logs/input/
logs/verify/
logs/legacy/
```
- NVIDIA 驱动、CUDA Wheel 与目标 GPU 的兼容性 主要事件:
- PaddleOCR-VL-1.6 GPU 模型初始化是否正常
- GPU 显存峰值和真实推理速度
- FP16/BF16、TensorRT 或批量推理收益
--- - `RUNTIME_PREPARED`
- `MODEL_INITIALIZED`
- `FILE_ROUTED`
- `PAGE_ROUTED`
- `PAGE_FINISHED`
- `TASK_COMPLETED`
- `IMAGE_COMPLETED`
- `DIRECTORY_SUMMARY`
- `VERIFY_COMPLETED`
### 优化总结 日志使用 UTF-8。Windows 控制台即使显示乱码,日志文件中的中文仍正常。
| 迭代 | 状态 | 结果 | ## 测试
|------|------|------|
| CPU 初始版本 | 已实测 | 后续单图约 238s |
| CPU `set_num_threads(20)` | 已实测 | 单图约 162s约 1.5x 加速 |
| CPU 多进程批量 | 已实现,稳定性依机器而定 | 理论提升批量吞吐;当前没有足够的可靠实测数据支持固定加速比 |
| 独立 GPU 子项目 | 已实现,未实机验证 | 等待 NVIDIA CUDA GPU 测试 |
> CPU 单图当前实测约 2.7 分钟。批量多进程主要提高总吞吐不会缩短某一张图片自身的推理延迟。GPU 性能在实机验证前不作预测。 ```bash
uv run --project cpu pytest -q
```
## 已知局限 当前测试覆盖:
| 问题 | 影响 | 说明 | - 页码范围解析
| ------------------ | ----------------------- | ------------------------------------- | - 文本标准化
| CPU 推理极慢 | 单图 ~2.7 min优化后 | 0.9B VL 模型不适合 CPU 实时场景 | - 有效文本层判定
| 自回归解码串行 | 无法更细粒度并行 | 生成阶段逐 token 依赖,多线程收益有限 | - 空文本/短文本自动回退条件
| 内存占用大 | 每进程需 ~2GB | 限制了 `batch_ocr.py` 并行度 | - 根入口设备解析
| Windows 控制台乱码 | 中文输出显示为乱码 | GBK 编码问题,文件写入/pipe 正常 |
| GPU 未实机验证 | 暂无 GPU 性能结论 | 当前机器只有集成显卡,需 NVIDIA CUDA GPU 验证 | 当前结果:
| ccache 警告 | 无实际影响 | 仅影响首次编译加速,可忽略 |
```text
17 passed
```
## 已验证状态
- CPU `verify`:通过
- 单图统一入口:假模型端到端通过,原真实模型能力此前已验证
- 批量统一入口:假模型端到端通过
- PDF `hybrid` 电子文本页:真实 PDF 验证通过,未加载模型
- PDF 扫描页自动回退 OCR假模型端到端通过
- PDF `text` / `ocr` 强制模式:通过
- GPU 环境路由:无环境时安全提示,不回退 CPU
- GPU 实机推理:尚未验证

View File

@ -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())

View File

@ -1,5 +1,29 @@
# Benchmark Results # 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/` 仅保留为可选的人工归档目录。

0
benchmarks/cpu/.gitkeep Normal file
View File

16
cpu/README.md Normal file
View File

@ -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
```

View File

@ -1,7 +1,7 @@
[project] [project]
name = "ocr-vl1-6" name = "ocr-vl1-6-cpu"
version = "0.1.0" version = "0.2.0"
description = "Add your description here" description = "CPU runtime for the unified PaddleOCR-VL-1.6 application"
readme = "README.md" readme = "README.md"
requires-python = ">=3.13" requires-python = ">=3.13"
dependencies = [ dependencies = [
@ -10,3 +10,8 @@ dependencies = [
"pypdfium2>=5.11.0", "pypdfium2>=5.11.0",
"setuptools>=83.0.0", "setuptools>=83.0.0",
] ]
[dependency-groups]
dev = [
"pytest>=8.4.0",
]

14
cpu/runner.py Normal file
View File

@ -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"))

View File

@ -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" }, { 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]] [[package]]
name = "jinja2" name = "jinja2"
version = "3.1.6" version = "3.1.6"
@ -1026,8 +1035,8 @@ wheels = [
] ]
[[package]] [[package]]
name = "ocr-vl1-6" name = "ocr-vl1-6-cpu"
version = "0.1.0" version = "0.2.0"
source = { virtual = "." } source = { virtual = "." }
dependencies = [ dependencies = [
{ name = "paddleocr", extra = ["doc-parser"] }, { name = "paddleocr", extra = ["doc-parser"] },
@ -1036,6 +1045,11 @@ dependencies = [
{ name = "setuptools" }, { name = "setuptools" },
] ]
[package.dev-dependencies]
dev = [
{ name = "pytest" },
]
[package.metadata] [package.metadata]
requires-dist = [ requires-dist = [
{ name = "paddleocr", extras = ["doc-parser"], specifier = "==3.7.0" }, { name = "paddleocr", extras = ["doc-parser"], specifier = "==3.7.0" },
@ -1044,6 +1058,9 @@ requires-dist = [
{ name = "setuptools", specifier = ">=83.0.0" }, { name = "setuptools", specifier = ">=83.0.0" },
] ]
[package.metadata.requires-dev]
dev = [{ name = "pytest", specifier = ">=8.4.0" }]
[[package]] [[package]]
name = "openai" name = "openai"
version = "2.45.0" 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" }, { 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]] [[package]]
name = "premailer" name = "premailer"
version = "3.10.0" 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" }, { 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]] [[package]]
name = "pypdfium2" name = "pypdfium2"
version = "5.11.0" 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" }, { 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]] [[package]]
name = "python-bidi" name = "python-bidi"
version = "0.6.11" version = "0.6.11"

View File

Before

Width:  |  Height:  |  Size: 31 KiB

After

Width:  |  Height:  |  Size: 31 KiB

View File

Before

Width:  |  Height:  |  Size: 206 KiB

After

Width:  |  Height:  |  Size: 206 KiB

View File

Before

Width:  |  Height:  |  Size: 329 KiB

After

Width:  |  Height:  |  Size: 329 KiB

View File

@ -1,144 +1,64 @@
# PaddleOCR-VL-1.6 GPU 子项目 # GPU 子项目
此目录是与根目录 CPU 版本隔离的 GPU 实验环境 GPU 环境与 CPU 环境隔离,所有用户功能通过仓库根目录统一入口调用
> **验证状态:未在 NVIDIA GPU 上实测。** 当前开发机器只有集成显卡,无法运行 CUDA。代码仅完成静态检查最终安装、兼容性和性能必须在目标 NVIDIA GPU 机器上验证。 > 当前开发机器没有 NVIDIA 独立显卡GPU 实机功能尚未验证。
## 为什么独立环境
`paddlepaddle``paddlepaddle-gpu` 都提供 `paddle` 模块,不能安全共用同一个虚拟环境。本目录具有独立的:
- `pyproject.toml`
- `.python-version`Python 3.11
- `.venv`(执行安装后生成)
- `uv.lock`(在目标 GPU 机器安装后生成)
根目录 CPU 环境不会被修改。
## 前置条件
1. NVIDIA CUDA GPUIntel/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 版本。
## 安装 ## 安装
在仓库根目录运行 依据 PaddlePaddle 官方兼容表选择 CUDA Wheel
```bash ```bash
# 只查看将执行的命令,不安装
python gpu/setup_env.py --cuda cu118 --dry-run python gpu/setup_env.py --cuda cu118 --dry-run
# CUDA 11.8 Wheel
python gpu/setup_env.py --cuda cu118 python gpu/setup_env.py --cuda cu118
```
# 或 CUDA 12.6 Wheel 或:
```bash
python gpu/setup_env.py --cuda cu126 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 ```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 构建 ## PDF 混合模式
- CUDA GPU 数量及名称
- `gpu:0` 是否能完成矩阵乘法
任何检查失败都会以非零状态退出,不会自动回退到 CPU。 `--pdf-mode hybrid` 会先读取 PDF 文本层。只有需要 OCR 的页面才初始化 GPU 模型,因此纯电子 PDF 不会创建 CUDA 模型或占用大块显存。
## 单图 Benchmark
```bash ```bash
uv run --project gpu python gpu/main.py python ocr.py data/documents/sample.pdf \
``` --device gpu \
常用参数:
```bash
uv run --project gpu python gpu/main.py \
--image images/手写01.png \
--device-id 0 \ --device-id 0 \
--warmup 1 \ --pdf-mode hybrid
--rounds 3
``` ```
Windows PowerShell 可写为单行,或使用反引号续行。 ## 当前限制
Benchmark 会记录: - 仅单 GPU
- 批量图片使用单模型串行处理
- 未启用 TensorRT/FP16/BF16
- 未验证 PaddleOCR-VL-1.6 在目标显卡上的显存需求
- 未实现多 GPU 调度
- 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`,按页渲染、逐页保存并支持断点续传:
```bash ```bash
uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf \ python ocr.py verify --device gpu
--device-id 0 \
--pages "1-10" \
--dpi 144
``` ```
常用选项:
```bash
# 中断后继续
uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf --resume
# 删除现有输出后重跑
uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf --overwrite
# 保留 PDF 页面的渲染 PNG
uv run --project gpu python gpu/pdf_ocr.py documents/sample.pdf --keep-rendered
```
无 CUDA 时脚本会立即退出,不会自动回落到 CPU。当前开发机器没有 NVIDIA GPU因此此入口尚未完成 GPU 实机验证。
PDF 日志默认写入:
```text
logs/pdf/<PDF名>-gpuN-YYYYMMDD-HHMMSS.log
```
日志与 `manifest.json` 会记录每页渲染、OCR、结果导出、状态保存、任务总用时和程序总用时。
## 当前范围
当前实现单 GPU、单图 Benchmark 和单 GPU PDF 逐页 OCR。暂未实现 GPU 多进程批处理,原因是:
- 同一 GPU 上启动多个模型实例会重复占用显存
- 多进程通常不会线性提升单卡吞吐
- 容易引发显存不足和 CUDA 上下文争抢
后续应优先评估模型/pipeline 原生批处理能力,再决定是否增加多 GPU 或任务队列。

View File

@ -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())

View File

@ -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())

View File

@ -1,7 +1,7 @@
[project] [project]
name = "ocr-vl1-6-gpu" name = "ocr-vl1-6-gpu"
version = "0.1.0" version = "0.2.0"
description = "GPU benchmark for PaddleOCR-VL-1.6" description = "GPU runtime for the unified PaddleOCR-VL-1.6 application"
readme = "README.md" readme = "README.md"
requires-python = ">=3.11,<3.13" requires-python = ">=3.11,<3.13"
dependencies = [ dependencies = [

14
gpu/runner.py Normal file
View File

@ -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"))

View File

@ -67,6 +67,10 @@ def main() -> int:
if args.dry_run: if args.dry_run:
return 0 return 0
ready_marker = project_dir / ".gpu-ready"
if ready_marker.exists():
ready_marker.unlink()
completed = subprocess.run(command, check=False) completed = subprocess.run(command, check=False)
if completed.returncode != 0: if completed.returncode != 0:
print( print(
@ -75,8 +79,12 @@ def main() -> int:
) )
return completed.returncode return completed.returncode
print("\n[OK] GPU 子项目环境已创建。下一步运行:") ready_marker.write_text(
print(f' uv run --project "{project_dir}" python "{project_dir / "verify_env.py"}"') f"cuda={args.cuda}\nindex={index_url}\n",
encoding="utf-8",
)
print("\n[OK] GPU 子项目环境已创建。下一步从仓库根目录运行:")
print(" python ocr.py verify --device gpu")
return 0 return 0

View File

@ -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())

140
main.py
View File

@ -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())

78
ocr.py Normal file
View File

@ -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())

3
ocr_app/__init__.py Normal file
View File

@ -0,0 +1,3 @@
"""Shared PaddleOCR-VL application package."""
__version__ = "0.2.0"

124
ocr_app/cli.py Normal file
View File

@ -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

497
ocr_app/commands.py Normal file
View File

@ -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

139
ocr_app/output.py Normal file
View File

@ -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

541
ocr_app/pdf.py Normal file
View File

@ -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()

127
ocr_app/pdf_text.py Normal file
View File

@ -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", "<EFBFBD>"} 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

170
ocr_app/runtime.py Normal file
View File

@ -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),
}

View File

@ -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())

View File

@ -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()

6
tests/conftest.py Normal file
View File

@ -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))

13
tests/test_launcher.py Normal file
View File

@ -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"

View File

@ -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
)

13
tests/test_page_spec.py Normal file
View File

@ -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)

102
tests/test_pdf_hybrid.py Normal file
View File

@ -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"

39
tests/test_pdf_text.py Normal file
View File

@ -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"