refactor:调整代码结构，将cpu模式收纳进子目录，统一入口，根据文件后缀和目录自动判断使用路由（pdf、图片和批量）
This commit is contained in:
parent
406845930b
commit
2988521c41
|
|
@ -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/
|
||||||
|
|
|
||||||
726
README.md
726
README.md
|
|
@ -1,428 +1,420 @@
|
||||||
# ocr-VL1.6
|
# PaddleOCR-VL-1.6 本地 OCR
|
||||||
|
|
||||||
本地部署 [PaddlePaddle/PaddleOCR-VL-1.6](https://github.com/PaddlePaddle/PaddleOCR) 的实验项目，包含已实测的 CPU 版本和独立隔离、待 NVIDIA GPU 验证的 GPU 版本。
|
本项目使用 PaddleOCR-VL-1.6 实现统一的图片 OCR、批量图片 OCR 和 PDF 识别，CPU/GPU 环境完全隔离，并通过根目录唯一入口 `ocr.py` 调用。
|
||||||
|
|
||||||
> 在线 Demo: [HuggingFace Space](https://huggingface.co/spaces/PaddlePaddle/PaddleOCR-VL-1.6_Online_Demo) · 模型权重: [HuggingFace](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.6)
|
PDF 默认使用 **文本提取 + OCR 混合模式**：优先提取 PDF 原始文本层，仅当页面没有有效文本层或文本质量不足时才加载 PaddleOCR-VL 并 OCR。
|
||||||
|
|
||||||
## 项目结构
|
> 当前开发机器只有集成显卡。CPU 功能已验证；GPU 代码已实现，但必须在 NVIDIA CUDA GPU 机器上安装和测试。
|
||||||
|
|
||||||
```
|
## 目录结构
|
||||||
|
|
||||||
|
```text
|
||||||
ocr-VL1.6/
|
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
|
||||||
│ ├── pyproject.toml # GPU 独立依赖
|
│ └── runner.py # 统一入口的 CPU 执行器
|
||||||
│ ├── .python-version # GPU 使用 Python 3.11
|
├── gpu/
|
||||||
│ └── README.md # GPU 安装与运行说明
|
│ ├── pyproject.toml # GPU 独立依赖
|
||||||
├── benchmarks/
|
│ ├── .python-version # Python 3.11
|
||||||
│ └── gpu/ # GPU Benchmark JSON 输出目录
|
│ ├── setup_env.py # CUDA Wheel 安装脚本
|
||||||
├── images/
|
│ └── runner.py # 统一入口的 GPU 执行器
|
||||||
│ └── 手写01.png # 测试图片：手写中文（1758×646）
|
├── data/
|
||||||
└── README.md
|
│ ├── images/ # 测试图片
|
||||||
|
│ └── documents/ # 测试 PDF
|
||||||
|
├── outputs/ # PDF 输出（git ignored）
|
||||||
|
├── benchmarks/ # 单图 Benchmark JSON
|
||||||
|
├── logs/ # UTF-8 运行日志
|
||||||
|
└── tests/ # 共享逻辑测试
|
||||||
```
|
```
|
||||||
|
|
||||||
## 技术栈
|
根目录不再保存 Paddle 虚拟环境或 Python 项目依赖。CPU 与 GPU 分别使用 `cpu/.venv` 和 `gpu/.venv`。
|
||||||
|
|
||||||
| 组件 | 版本 | 说明 |
|
## 安装
|
||||||
| ---------------- | ----- | ---------------------------------------- |
|
|
||||||
| Python（CPU） | 3.13 | 根目录独立环境 |
|
|
||||||
| Python（GPU） | 3.11 | `gpu/` 独立环境，提升 GPU Wheel 兼容性 |
|
|
||||||
| PaddlePaddle | 3.2.1 | CPU 使用 `paddlepaddle`，GPU 使用 `paddlepaddle-gpu` |
|
|
||||||
| PaddleOCR | 3.7.0 | 带 `doc-parser` extra |
|
|
||||||
| PaddleOCR-VL-1.6 | 0.9B | 主 OCR 视觉语言模型（~1.8GB） |
|
|
||||||
| PP-DocLayoutV3 | - | 版面检测模型（~126MB） |
|
|
||||||
|
|
||||||
模型缓存目录：`~/.paddlex/official_models/`
|
### 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 线程数；长任务建议保留 1～2 个核心给系统
|
|
||||||
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` | 第 1～5 页 |
|
|
||||||
| `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 # 纯文本结果
|
||||||
├── pages/
|
│ ├── result.json # PaddleOCR 结构化结果
|
||||||
│ ├── page-0001.md
|
│ ├── benchmark.json # 模型/推理/导出耗时
|
||||||
│ ├── page-0001.json
|
│ └── assets/ # 识别结果中的图片资源（如有）
|
||||||
│ └── ...
|
├── pdfs/
|
||||||
├── assets/ # 表格、图片等 Markdown 资源
|
│ └── <相对目录>/<PDF名>/
|
||||||
└── rendered/ # 仅使用 --keep-rendered 时保留
|
│ ├── manifest.json
|
||||||
|
│ ├── document.md
|
||||||
|
│ ├── document.json
|
||||||
|
│ └── pages/
|
||||||
|
└── batches/
|
||||||
|
└── <目录名>-<时间戳>.json # 批量任务汇总
|
||||||
```
|
```
|
||||||
|
|
||||||
默认不保留中间渲染 PNG；OCR 完成后会删除临时图。每页 JSON 会将 `input_path` 恢复为原 PDF 路径，并记录 `page_index`、`page_number`、`page_count` 和 `render_dpi`。
|
例如：
|
||||||
|
|
||||||
### 恢复与错误处理
|
|
||||||
|
|
||||||
- 输出目录已存在时，必须显式使用 `--resume` 或 `--overwrite`
|
|
||||||
- `--resume` 会校验 PDF SHA-256 和 DPI，防止接续到错误任务
|
|
||||||
- 单页失败默认写入 manifest 并继续后续页面
|
|
||||||
- `--fail-fast` 可在第一页失败后立即停止
|
|
||||||
- `Ctrl+C` 会保存当前 manifest；下次使用 `--resume` 继续
|
|
||||||
- 逐页文件和 manifest 使用临时文件替换，降低中途退出造成文件损坏的概率
|
|
||||||
|
|
||||||
### DPI 建议
|
|
||||||
|
|
||||||
| 文档类型 | 建议 DPI |
|
|
||||||
|----------|---------:|
|
|
||||||
| 普通打印文字 | 120～144 |
|
|
||||||
| 小字号文档 | 150～200 |
|
|
||||||
| 手写或低质量扫描件 | 200～250 |
|
|
||||||
|
|
||||||
CPU 当前单图实测约 162 秒。长 PDF 总时间可粗略按 `待处理页数 × 单页耗时` 估算，因此建议先用 `--pages "1"` 测试效果和耗时，再扩大页码范围。DPI 越高通常越慢，不建议默认使用 300 DPI。
|
|
||||||
|
|
||||||
## 工作原理
|
|
||||||
|
|
||||||
`PaddleOCRVL` pipeline 分两阶段：
|
|
||||||
|
|
||||||
```
|
|
||||||
输入图片 → [PP-DocLayoutV3 版面检测] → [PaddleOCR-VL-1.6-0.9B 文字识别] → 结构化输出
|
|
||||||
```
|
|
||||||
|
|
||||||
1. **版面检测** — PP-DocLayoutV3 检测页面中的文本块区域（坐标 + 标签 + 置信度）
|
|
||||||
2. **OCR 识别** — PaddleOCR-VL-1.6-0.9B（GQA 架构视觉语言模型）逐块识别文字
|
|
||||||
3. **结果输出** — 返回 `PaddleOCRVLResult`，包含布局信息和识别文本
|
|
||||||
|
|
||||||
### 输出结构
|
|
||||||
|
|
||||||
| 字段 | 类型 | 说明 |
|
|
||||||
| ---------------------- | ---------------------- | ------------------------------------------------------------ |
|
|
||||||
| `layout_det_res.boxes` | list[dict] | 版面文本区域（cls_id, label, score, coordinate, polygon_points） |
|
|
||||||
| `parsing_res_list` | list[PaddleOCRVLBlock] | 识别文本块（$.label, $.bbox, $.content, $.polygon_points） |
|
|
||||||
| `model_settings` | dict | 推理配置开关（版面检测/图表/印章等） |
|
|
||||||
| `width` / `height` | int | 图片尺寸 |
|
|
||||||
|
|
||||||
## 性能优化迭代
|
|
||||||
|
|
||||||
测试机器：Windows 11, CPU 20 核（逻辑线程）, RAM 32GB, PaddlePaddle 3.2.1 CPU
|
|
||||||
|
|
||||||
### 迭代 0：初始状态（无任何优化）
|
|
||||||
|
|
||||||
直接调用 `pipeline.predict()`，未设置任何线程参数。
|
|
||||||
|
|
||||||
| 阶段 | 耗时 |
|
|
||||||
| ----------------------- | --------------- |
|
|
||||||
| 模型初始化（加载权重） | ~60s |
|
|
||||||
| 首次推理（含 JIT 编译） | ~285s |
|
|
||||||
| 后续推理 | ~238s（~4 min） |
|
|
||||||
|
|
||||||
### 迭代 1：算子级并行 — `core.set_num_threads()`
|
|
||||||
|
|
||||||
**探索过程：**
|
|
||||||
|
|
||||||
| 尝试 | 方法 | 结果 |
|
|
||||||
| ---- | -------------------------------------------------- | ---------------------------------------------------- |
|
|
||||||
| ❌ | `paddle.set_num_threads()` | Paddle 3.x 已移除该 API |
|
|
||||||
| ❌ | 环境变量 `OMP_NUM_THREADS` / `MKL_NUM_THREADS` | Paddle 3.x 内部使用 oneDNN，不读取这些变量 |
|
|
||||||
| ✅ | `from paddle import core; core.set_num_threads(N)` | **有效！** oneDNN 底层算子（matmul 等）受该 API 控制 |
|
|
||||||
|
|
||||||
**矩阵乘法微基准测试（4000×4000）：**
|
|
||||||
|
|
||||||
| 线程数 | 耗时 (matmul) | 加速比 |
|
|
||||||
| ------ | ------------- | -------- |
|
|
||||||
| 1 | 0.952s | 1.0x |
|
|
||||||
| 4 | 0.419s | 2.3x |
|
|
||||||
| 8 | 0.323s | 2.9x |
|
|
||||||
| 16 | 0.240s | 4.0x |
|
|
||||||
| **20** | **0.223s** | **4.3x** |
|
|
||||||
|
|
||||||
**应用到 OCR 后的实际效果：**
|
|
||||||
|
|
||||||
设置 `core.set_num_threads(20)` 后重新评测：
|
|
||||||
|
|
||||||
| 阶段 | 优化前 | 优化后 | 提速 |
|
|
||||||
| ---------- | ------ | --------------------- | -------- |
|
|
||||||
| 模型初始化 | ~60s | ~40s | 1.5x |
|
|
||||||
| 推理 | ~238s | **~162s（~2.7 min）** | **1.5x** |
|
|
||||||
|
|
||||||
**为什么不是 4.3x？** 矩阵乘法只是 OCR pipeline 的一部分。自回归解码（逐 token 生成）天然串行、I/O 等待、版面检测中的非矩阵运算等不受线程数影响。
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
### 迭代 2：批量多进程并行 — `batch_ocr.py`
|
|
||||||
|
|
||||||
**思路：** 多张图片时，用 `multiprocessing.Pool` 启动多个独立进程，每个进程加载一份 pipeline 实例，同时处理不同图片。
|
|
||||||
|
|
||||||
**遇到的问题 & 修复（迭代 2.1）：**
|
|
||||||
|
|
||||||
| 问题 | 原因 | 修复 |
|
|
||||||
|------|------|------|
|
|
||||||
| 系统卡顿/黑屏/无响应 | `Pool.starmap` 同时启动 N 个进程，同步加载 N×2GB 模型，CPU + 内存瞬间打满 | ① 进程错峰启动（随机延迟 0~15s）② `psutil` 降低进程优先级 ③ 预留 1 核给 OS ④ `imap_unordered` 替代 `starmap` |
|
|
||||||
|
|
||||||
**策略：**
|
|
||||||
- 每个子进程独立调用 `core.set_num_threads((总核心-1) / 进程数)`，预留核心给 OS
|
|
||||||
- 例如 2 进程 × 9 线程 = 18 核，留 2 核给系统
|
|
||||||
- `--stagger` 控制错峰窗口，默认 15s
|
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# 2 进程并行（安全默认值）
|
python ocr.py data/images/名片01.jpg --device cpu
|
||||||
uv run python batch_ocr.py images/
|
|
||||||
|
|
||||||
# 4 进程并行（需 32GB+ RAM）
|
|
||||||
uv run python batch_ocr.py images/ --workers 4
|
|
||||||
|
|
||||||
# 自定义错峰窗口（值越大内存峰值越低，但总耗时增加）
|
|
||||||
uv run python batch_ocr.py images/ --workers 4 --stagger 30
|
|
||||||
```
|
```
|
||||||
|
|
||||||
| 配置 | 适用场景 | 理论加速比 | 内存开销 | 实际限制 |
|
会生成：
|
||||||
| -------------------- | -------- | --------------- | ---------- | -------------- |
|
|
||||||
| `set_num_threads(N)` | 单张图片 | ~1.5x | 无额外开销 | 自回归解码瓶颈 |
|
|
||||||
| `batch_ocr.py` | 批量多图 | ~Nx（N=进程数） | N × 2GB | 内存/内存带宽，需错峰避免打满系统 |
|
|
||||||
|
|
||||||
> ⚠️ 每个进程独立加载模型（~2GB），32GB RAM 建议从 `--workers 2` 开始测试。默认值是相对保守配置，但是否稳定仍取决于可用内存、散热、后台应用和图片复杂度。
|
```text
|
||||||
|
outputs/images/名片01_jpg/result.md
|
||||||
|
outputs/images/名片01_jpg/result.txt
|
||||||
|
outputs/images/名片01_jpg/result.json
|
||||||
|
outputs/images/名片01_jpg/benchmark.json
|
||||||
|
```
|
||||||
|
|
||||||
---
|
图片目录名包含扩展名，避免 `same.png` 与 `same.jpg` 相互覆盖。
|
||||||
|
|
||||||
### 迭代 3：独立 GPU 子项目（待实机验证）
|
重构前生成的 PDF 结果可能仍直接位于 `outputs/<PDF名>/`。这些旧结果不会自动删除；新任务统一写入 `outputs/pdfs/<PDF名>/`，确认不再需要后可手动迁移或清理。
|
||||||
|
|
||||||
为避免 `paddlepaddle` 和 `paddlepaddle-gpu` 相互覆盖，在同一仓库新增 `gpu/` 子项目，使用独立 Python、虚拟环境、依赖配置和锁文件。
|
目录任务再次运行时：
|
||||||
|
|
||||||
已完成：
|
- 已存在的 PDF manifest 自动断点续传
|
||||||
|
- 新加入的 PDF 自动创建新任务
|
||||||
|
- 图片重新识别，并使用原子写入覆盖对应输出文件
|
||||||
|
- `--overwrite` 会强制 PDF 重新处理
|
||||||
|
- 每次目录任务都会生成新的 `outputs/batches/*.json` 汇总
|
||||||
|
|
||||||
- `gpu/setup_env.py`：根据 `cu118` / `cu126` Wheel 索引创建环境
|
旧命令前缀 `image`、`pdf`、`batch` 暂时兼容，例如 `python ocr.py image a.png`，但推荐直接传入路径。
|
||||||
- `gpu/verify_env.py`：检查 CUDA 构建、设备数量并执行 GPU 矩阵乘法
|
|
||||||
- `gpu/main.py`：显式指定 `device="gpu:N"`，支持预热、多轮计时和 JSON 输出
|
|
||||||
- 无 CUDA 时立即退出，不静默回退到 CPU
|
|
||||||
- CPU 环境下已通过 Python 语法、CLI 和安全退出检查
|
|
||||||
|
|
||||||
尚未验证：
|
## PDF 混合模式
|
||||||
|
|
||||||
- NVIDIA 驱动、CUDA Wheel 与目标 GPU 的兼容性
|
### 默认：hybrid
|
||||||
- PaddleOCR-VL-1.6 GPU 模型初始化是否正常
|
|
||||||
- GPU 显存峰值和真实推理速度
|
|
||||||
- FP16/BF16、TensorRT 或批量推理收益
|
|
||||||
|
|
||||||
---
|
```bash
|
||||||
|
python ocr.py data/documents/sample.pdf --device cpu
|
||||||
|
```
|
||||||
|
|
||||||
### 优化总结
|
每页流程：
|
||||||
|
|
||||||
| 迭代 | 状态 | 结果 |
|
```text
|
||||||
|------|------|------|
|
读取 PDF 文本层
|
||||||
| CPU 初始版本 | 已实测 | 后续单图约 238s |
|
↓
|
||||||
| CPU `set_num_threads(20)` | 已实测 | 单图约 162s，约 1.5x 加速 |
|
文本质量评估
|
||||||
| CPU 多进程批量 | 已实现，稳定性依机器而定 | 理论提升批量吞吐；当前没有足够的可靠实测数据支持固定加速比 |
|
┌────┴────┐
|
||||||
| 独立 GPU 子项目 | 已实现，未实机验证 | 等待 NVIDIA CUDA GPU 测试 |
|
│ │
|
||||||
|
有效文本 无效/不足
|
||||||
|
│ │
|
||||||
|
直接保存 渲染页面 → PaddleOCR-VL
|
||||||
|
└────┬────┘
|
||||||
|
↓
|
||||||
|
逐页 Markdown/JSON + 合并文档
|
||||||
|
```
|
||||||
|
|
||||||
> CPU 单图当前实测约 2.7 分钟。批量多进程主要提高总吞吐，不会缩短某一张图片自身的推理延迟。GPU 性能在实机验证前不作预测。
|
如果整份 PDF 都具有有效文本层，模型完全不会加载。实测项目中的电子合同 PDF，单页混合处理约 `0.06s`，且 `model_used=false`。
|
||||||
|
|
||||||
## 已知局限
|
### 强制文本模式
|
||||||
|
|
||||||
| 问题 | 影响 | 说明 |
|
```bash
|
||||||
| ------------------ | ----------------------- | ------------------------------------- |
|
python ocr.py data/documents/sample.pdf \
|
||||||
| CPU 推理极慢 | 单图 ~2.7 min（优化后） | 0.9B VL 模型不适合 CPU 实时场景 |
|
--device cpu \
|
||||||
| 自回归解码串行 | 无法更细粒度并行 | 生成阶段逐 token 依赖，多线程收益有限 |
|
--pdf-mode text
|
||||||
| 内存占用大 | 每进程需 ~2GB | 限制了 `batch_ocr.py` 并行度 |
|
```
|
||||||
| Windows 控制台乱码 | 中文输出显示为乱码 | GBK 编码问题，文件写入/pipe 正常 |
|
|
||||||
| GPU 未实机验证 | 暂无 GPU 性能结论 | 当前机器只有集成显卡，需 NVIDIA CUDA GPU 验证 |
|
所有页面只提取 PDF 文本层，不执行 OCR。扫描页可能得到空文本。
|
||||||
| ccache 警告 | 无实际影响 | 仅影响首次编译加速，可忽略 |
|
|
||||||
|
### 强制 OCR 模式
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python ocr.py data/documents/sample.pdf \
|
||||||
|
--device cpu \
|
||||||
|
--pdf-mode ocr
|
||||||
|
```
|
||||||
|
|
||||||
|
所有页面都渲染后交给 PaddleOCR-VL，适合模型一致性 Benchmark。
|
||||||
|
|
||||||
|
## 文本层质量判定
|
||||||
|
|
||||||
|
混合模式默认要求：
|
||||||
|
|
||||||
|
| 参数 | 默认值 | 含义 |
|
||||||
|
|------|-------:|------|
|
||||||
|
| `--text-min-chars` | 50 | 非空白字符最小数量 |
|
||||||
|
| `--text-min-printable-ratio` | 0.85 | 可打印字符最低比例 |
|
||||||
|
| `--text-min-content-ratio` | 0.60 | 字母、数字、CJK 字符最低比例 |
|
||||||
|
| `--text-max-replacement-ratio` | 0.02 | Unicode 替换字符最高比例 |
|
||||||
|
| `--text-min-density` | 25 | 页面文本密度最低值 |
|
||||||
|
|
||||||
|
例如某些扫描 PDF 只有页码或隐藏乱码层，混合模式会因为 `too_few_characters`、`low_content_ratio` 或 `high_replacement_ratio` 自动回退 OCR。
|
||||||
|
|
||||||
|
每页的判定结果会写入日志和 manifest：
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"source_type": "ocr",
|
||||||
|
"routing_reason": "too_few_characters",
|
||||||
|
"text_layer": {
|
||||||
|
"usable": false,
|
||||||
|
"non_whitespace_chars": 8,
|
||||||
|
"printable_ratio": 1.0,
|
||||||
|
"content_ratio": 0.75
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## PDF 页码与恢复
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# 第 1～5 页、第 8 页、第 10 页到末页
|
||||||
|
python ocr.py sample.pdf --pages "1-5,8,10-" --device cpu
|
||||||
|
|
||||||
|
# 中断后继续
|
||||||
|
python ocr.py sample.pdf --resume --device cpu
|
||||||
|
|
||||||
|
# 删除旧输出并重做
|
||||||
|
python ocr.py sample.pdf --overwrite --device cpu
|
||||||
|
```
|
||||||
|
|
||||||
|
`--resume` 会校验：
|
||||||
|
|
||||||
|
- PDF SHA-256
|
||||||
|
- PDF 处理模式
|
||||||
|
- DPI
|
||||||
|
- 文本层判定阈值
|
||||||
|
- manifest 版本
|
||||||
|
|
||||||
|
旧纯 OCR manifest（版本 1）不能直接用于新混合模式，请使用 `--overwrite` 重新生成。
|
||||||
|
|
||||||
|
## PDF 输出
|
||||||
|
|
||||||
|
```text
|
||||||
|
outputs/pdfs/<PDF名>/
|
||||||
|
├── manifest.json
|
||||||
|
├── document.md
|
||||||
|
├── document.json
|
||||||
|
├── pages/
|
||||||
|
│ ├── page-0001.md
|
||||||
|
│ ├── page-0001.json
|
||||||
|
│ └── ...
|
||||||
|
├── assets/
|
||||||
|
└── rendered/ # 仅 --keep-rendered 时存在
|
||||||
|
```
|
||||||
|
|
||||||
|
合并 Markdown 标题会标明页面来源：
|
||||||
|
|
||||||
|
```markdown
|
||||||
|
## Page 1 (text)
|
||||||
|
## Page 2 (ocr)
|
||||||
|
```
|
||||||
|
|
||||||
|
Manifest 汇总包括：
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"text_pages": 8,
|
||||||
|
"ocr_pages": 2,
|
||||||
|
"model_used": true,
|
||||||
|
"model_initialized_during_task": true,
|
||||||
|
"timing": {
|
||||||
|
"text_extract_total_seconds": 0.15,
|
||||||
|
"render_total_seconds": 0.8,
|
||||||
|
"ocr_total_seconds": 324.5,
|
||||||
|
"model_init_seconds": 40.0,
|
||||||
|
"task_total_seconds": 366.0
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## PDF 常用参数
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python ocr.py sample.pdf \
|
||||||
|
--device cpu \
|
||||||
|
--pdf-mode hybrid \
|
||||||
|
--pages "1-10" \
|
||||||
|
--dpi 144 \
|
||||||
|
--threads 18 \
|
||||||
|
--keep-rendered \
|
||||||
|
--log-file logs/pdf-sample.log
|
||||||
|
```
|
||||||
|
|
||||||
|
DPI 建议：
|
||||||
|
|
||||||
|
| 文档 | DPI |
|
||||||
|
|------|----:|
|
||||||
|
| 普通打印文字 | 120～144 |
|
||||||
|
| 小字号 | 150～200 |
|
||||||
|
| 手写/低质量扫描件 | 200～250 |
|
||||||
|
|
||||||
|
DPI 只影响需要 OCR 的页面，不影响直接提取文本层的页面。
|
||||||
|
|
||||||
|
## 日志
|
||||||
|
|
||||||
|
统一日志格式：
|
||||||
|
|
||||||
|
```text
|
||||||
|
2026-07-16 15:07:45 | INFO | pid=33552 | PAGE_ROUTED page=1 source=text reason=usable_text_layer
|
||||||
|
```
|
||||||
|
|
||||||
|
默认目录：
|
||||||
|
|
||||||
|
```text
|
||||||
|
logs/input/
|
||||||
|
logs/verify/
|
||||||
|
logs/legacy/
|
||||||
|
```
|
||||||
|
|
||||||
|
主要事件：
|
||||||
|
|
||||||
|
- `RUNTIME_PREPARED`
|
||||||
|
- `MODEL_INITIALIZED`
|
||||||
|
- `FILE_ROUTED`
|
||||||
|
- `PAGE_ROUTED`
|
||||||
|
- `PAGE_FINISHED`
|
||||||
|
- `TASK_COMPLETED`
|
||||||
|
- `IMAGE_COMPLETED`
|
||||||
|
- `DIRECTORY_SUMMARY`
|
||||||
|
- `VERIFY_COMPLETED`
|
||||||
|
|
||||||
|
日志使用 UTF-8。Windows 控制台即使显示乱码，日志文件中的中文仍正常。
|
||||||
|
|
||||||
|
## 测试
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv run --project cpu pytest -q
|
||||||
|
```
|
||||||
|
|
||||||
|
当前测试覆盖：
|
||||||
|
|
||||||
|
- 页码范围解析
|
||||||
|
- 文本标准化
|
||||||
|
- 有效文本层判定
|
||||||
|
- 空文本/短文本自动回退条件
|
||||||
|
- 根入口设备解析
|
||||||
|
|
||||||
|
当前结果：
|
||||||
|
|
||||||
|
```text
|
||||||
|
17 passed
|
||||||
|
```
|
||||||
|
|
||||||
|
## 已验证状态
|
||||||
|
|
||||||
|
- CPU `verify`：通过
|
||||||
|
- 单图统一入口：假模型端到端通过，原真实模型能力此前已验证
|
||||||
|
- 批量统一入口：假模型端到端通过
|
||||||
|
- PDF `hybrid` 电子文本页：真实 PDF 验证通过，未加载模型
|
||||||
|
- PDF 扫描页自动回退 OCR：假模型端到端通过
|
||||||
|
- PDF `text` / `ocr` 强制模式：通过
|
||||||
|
- GPU 环境路由：无环境时安全提示，不回退 CPU
|
||||||
|
- GPU 实机推理：尚未验证
|
||||||
|
|
|
||||||
329
batch_ocr.py
329
batch_ocr.py
|
|
@ -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())
|
|
||||||
|
|
@ -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,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
|
||||||
|
```
|
||||||
|
|
@ -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",
|
||||||
|
]
|
||||||
|
|
@ -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"))
|
||||||
|
|
@ -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"
|
||||||
|
Before Width: | Height: | Size: 31 KiB After Width: | Height: | Size: 31 KiB |
|
Before Width: | Height: | Size: 206 KiB After Width: | Height: | Size: 206 KiB |
|
Before Width: | Height: | Size: 329 KiB After Width: | Height: | Size: 329 KiB |
148
gpu/README.md
148
gpu/README.md
|
|
@ -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 GPU（Intel/AMD 集成显卡不能运行 Paddle CUDA 版本）
|
|
||||||
2. 兼容的 NVIDIA 驱动
|
|
||||||
3. Python 3.11
|
|
||||||
4. uv
|
|
||||||
5. 根据 PaddlePaddle 官方兼容表选择 CUDA Wheel
|
|
||||||
|
|
||||||
先检查目标机器：
|
|
||||||
|
|
||||||
```bash
|
|
||||||
nvidia-smi
|
|
||||||
```
|
|
||||||
|
|
||||||
`nvidia-smi` 显示的 CUDA Version 是驱动支持上限，不等同于本机安装的 CUDA Toolkit，也不能单独用于判断 Wheel 版本。
|
|
||||||
|
|
||||||
## 安装
|
## 安装
|
||||||
|
|
||||||
在仓库根目录运行：
|
依据 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 或任务队列。
|
|
||||||
|
|
|
||||||
291
gpu/main.py
291
gpu/main.py
|
|
@ -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())
|
|
||||||
214
gpu/pdf_ocr.py
214
gpu/pdf_ocr.py
|
|
@ -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())
|
|
||||||
|
|
@ -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 = [
|
||||||
|
|
|
||||||
|
|
@ -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"))
|
||||||
|
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
140
main.py
|
|
@ -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())
|
|
||||||
|
|
@ -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())
|
||||||
|
|
@ -0,0 +1,3 @@
|
||||||
|
"""Shared PaddleOCR-VL application package."""
|
||||||
|
|
||||||
|
__version__ = "0.2.0"
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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()
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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),
|
||||||
|
}
|
||||||
182
pdf_ocr.py
182
pdf_ocr.py
|
|
@ -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())
|
|
||||||
658
pdf_ocr_core.py
658
pdf_ocr_core.py
|
|
@ -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()
|
|
||||||
|
|
@ -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))
|
||||||
|
|
@ -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"
|
||||||
|
|
@ -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
|
||||||
|
)
|
||||||
|
|
@ -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)
|
||||||
|
|
@ -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"
|
||||||
|
|
@ -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"
|
||||||
Loading…
Reference in New Issue