feat:仿照VL1.6实现,并添加了模型选择,默认是medium模型,还有small、tiny可以选

This commit is contained in:
kuuhaku 2026-07-17 10:19:05 +08:00
commit 78bd690b7e
42 changed files with 5068 additions and 0 deletions

21
.gitignore vendored Normal file
View File

@ -0,0 +1,21 @@
__pycache__/
*.py[oc]
build/
dist/
wheels/
*.egg-info/
.pytest_cache/
.venv/
cpu/.venv/
gpu/.venv/
gpu/.gpu-ready
outputs/*
!outputs/.gitkeep
logs/*
!logs/.gitkeep
benchmarks/cpu/*.json
benchmarks/gpu/*.json
!benchmarks/cpu/.gitkeep
!benchmarks/gpu/.gitkeep

196
README.md Normal file
View File

@ -0,0 +1,196 @@
# PP-OCRv6 本地 OCR
本项目使用 PP-OCRv6 实现图片、PDF 和目录批量 OCR。工程结构参考 `../ocr-VL1.6`CPU/GPU 环境隔离,并通过根目录 `ocr.py` 统一调用。
输出 Markdown 的定位是**便于阅读的文字结果**。本项目不实现版面分析、表格结构恢复或富文档还原。
## 功能
- 单图片 PP-OCRv6 识别
- 图片与 PDF 目录批处理
- PDF `hybrid` / `text` / `ocr` 模式
- PDF 文本层质量判断和扫描页自动 OCR
- PDF 页码选择、断点续传与覆盖重做
- TXT、简单 Markdown、归一化 JSON、Benchmark
- 可选 PaddleOCR 原始 JSON 和 OCR 可视化
- CPU/GPU 隔离GPU 失败不回退 CPU
- 模型延迟加载,同一批次复用模型
支持三种成对的检测与识别模型规格:
| 参数 | 检测模型 | 识别模型 |
|---|---|---|
| `tiny` | `PP-OCRv6_tiny_det` | `PP-OCRv6_tiny_rec` |
| `small` | `PP-OCRv6_small_det` | `PP-OCRv6_small_rec` |
| `medium` | `PP-OCRv6_medium_det` | `PP-OCRv6_medium_rec` |
不传参数时默认使用 `medium`
## 安装
### CPU
```bash
uv sync --project cpu
```
### GPU
```bash
python gpu/setup_env.py --cuda cu118 --dry-run
python gpu/setup_env.py --cuda cu118
```
也可按目标机器使用 `cu126`。CUDA、驱动和 PaddlePaddle Wheel 必须相互兼容。
## 使用
### 验证环境
```bash
python ocr.py verify --device cpu
python ocr.py verify --device gpu
```
`verify` 只验证 Paddle 运行环境,不下载和初始化 OCR 模型。
### 单图片
```bash
# 默认 medium
python ocr.py data/images/手写01.png --device cpu
# 选择 small
python ocr.py data/images/手写01.png --device cpu --model-size small
# 选择 tiny--model 是 --model-size 的别名
python ocr.py data/images/手写01.png --device cpu --model tiny
```
保存可视化和原始结果:
```bash
python ocr.py data/images/手写01.png \
--device cpu \
--save-visualization \
--save-raw-result
```
多轮 Benchmark
```bash
python ocr.py data/images/手写01.png --warmup 1 --rounds 3
```
### PDF
```bash
# 默认混合模式:优先文本层,扫描页才 OCR
python ocr.py data/documents/sample.pdf --device cpu
# 强制文本层
python ocr.py sample.pdf --pdf-mode text
# 强制全部页面 OCR
python ocr.py sample.pdf --pdf-mode ocr
# 指定页码
python ocr.py sample.pdf --pages "1-5,8,10-"
# 中断后恢复
python ocr.py sample.pdf --resume
# 删除原任务并重做
python ocr.py sample.pdf --overwrite
```
### 目录
```bash
python ocr.py data/ --device cpu
python ocr.py data/ --recursive --device cpu
```
目录任务串行处理文件,避免无控制并发造成内存占用和界面卡顿。
## 主要 PP-OCRv6 参数
```text
--model-size tiny|small|medium
--model tiny|small|medium # --model-size 的别名
--lang ch
--text-rec-score-thresh 0.0
--text-det-limit-side-len N
--text-det-limit-type min|max
--text-det-thresh FLOAT
--text-det-box-thresh FLOAT
--text-det-unclip-ratio FLOAT
--text-recognition-batch-size 6
--return-word-box
--doc-orientation-classify / --no-doc-orientation-classify
--doc-unwarping / --no-doc-unwarping
--textline-orientation / --no-textline-orientation
```
默认开启文档方向分类和文本行方向分类,默认关闭文档去畸变。
## 输出
```text
outputs/
├── images/
│ └── <图片名_扩展名>/
│ ├── result.txt
│ ├── result.md
│ ├── result.json
│ ├── benchmark.json
│ ├── raw-result.json # 可选
│ └── visualization.jpg # 可选
├── pdfs/
│ └── <PDF名>/
│ ├── manifest.json
│ ├── document.md
│ ├── document.json
│ └── pages/
└── batches/
```
`result.json` 使用项目自有稳定结构,主要包含:
- 文本行
- 识别置信度
- 多边形与矩形坐标
- 模型和语言信息
- 图片尺寸
- 汇总统计
## PDF 混合模式
每页先提取 PDF 文本层:
```text
有效文本层 → 直接保存文本
无效文本层 → 渲染 PNG → PP-OCRv6
```
纯电子 PDF 不加载 PP-OCRv6 模型。Manifest 会记录 PDF 哈希、DPI、文本层阈值和模型配置关键配置变化时必须使用 `--overwrite`,不能错误续传。
## 能力边界
本项目只提供通用文字检测和识别:
- Markdown 是按识别顺序拼接的便读文本;
- 不恢复标题层级;
- 不恢复多栏版面;
- 不恢复表格单元格结构;
- 不提供与 PaddleOCR-VL 等价的富 Markdown。
JSON 中保留坐标,调用方可以按具体业务继续处理。
## 测试
```bash
uv run --project cpu pytest -q
```
测试使用假模型覆盖结果适配、输出路由和 PDF 混合流程,不要求每次测试都下载 OCR 模型。

19
benchmarks/README.md Normal file
View File

@ -0,0 +1,19 @@
# Benchmark
图片 Benchmark 默认保存在:
```text
outputs/images/<图片名_扩展名>/benchmark.json
```
示例:
```bash
python ocr.py data/images/手写01.png --warmup 1 --rounds 3
```
需要额外复制时:
```bash
python ocr.py data/images/手写01.png --benchmark-json benchmarks/cpu/手写01.json
```

0
benchmarks/cpu/.gitkeep Normal file
View File

0
benchmarks/gpu/.gitkeep Normal file
View File

1
cpu/.python-version Normal file
View File

@ -0,0 +1 @@
3.13

14
cpu/README.md Normal file
View File

@ -0,0 +1,14 @@
# CPU 子项目
安装:
```bash
uv sync --project cpu
```
统一从仓库根目录运行:
```bash
python ocr.py verify --device cpu
python ocr.py data/images/手写01.png --device cpu
```

17
cpu/pyproject.toml Normal file
View File

@ -0,0 +1,17 @@
[project]
name = "pp-ocrv6-cpu"
version = "0.1.0"
description = "CPU runtime for the unified PP-OCRv6 application"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"paddleocr==3.7.0",
"paddlepaddle==3.2.1",
"pypdfium2>=5.11.0",
"setuptools>=83.0.0",
]
[dependency-groups]
dev = [
"pytest>=8.4.0",
]

14
cpu/runner.py Normal file
View File

@ -0,0 +1,14 @@
"""CPU environment runner used by the root unified launcher."""
from pathlib import Path
import sys
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from ocr_app.cli import main
if __name__ == "__main__":
raise SystemExit(main(device="cpu"))

1656
cpu/uv.lock Normal file

File diff suppressed because it is too large Load Diff

BIN
data/images/合同01.jpg Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 214 KiB

BIN
data/images/名片01.jpg Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 31 KiB

BIN
data/images/名片02.jpg Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 206 KiB

BIN
data/images/手写01.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 329 KiB

File diff suppressed because it is too large Load Diff

1
gpu/.python-version Normal file
View File

@ -0,0 +1 @@
3.11

17
gpu/README.md Normal file
View File

@ -0,0 +1,17 @@
# GPU 子项目
依据目标 CUDA 版本安装:
```bash
python gpu/setup_env.py --cuda cu118
# 或
python gpu/setup_env.py --cuda cu126
```
安装后验证:
```bash
python ocr.py verify --device gpu
```
GPU 失败时不会自动回退 CPU。

12
gpu/pyproject.toml Normal file
View File

@ -0,0 +1,12 @@
[project]
name = "pp-ocrv6-gpu"
version = "0.1.0"
description = "GPU runtime for the unified PP-OCRv6 application"
readme = "README.md"
requires-python = ">=3.11,<3.13"
dependencies = [
"paddleocr==3.7.0",
"paddlepaddle-gpu==3.2.1",
"pypdfium2>=5.11.0",
"setuptools>=83.0.0",
]

14
gpu/runner.py Normal file
View File

@ -0,0 +1,14 @@
"""GPU environment runner used by the root unified launcher."""
from pathlib import Path
import sys
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from ocr_app.cli import main
if __name__ == "__main__":
raise SystemExit(main(device="gpu"))

92
gpu/setup_env.py Normal file
View File

@ -0,0 +1,92 @@
"""根据目标 CUDA 版本创建独立的 GPU uv 环境。"""
import argparse
import shutil
import subprocess
import sys
from pathlib import Path
PADDLE_INDEXES = {
"cu118": "https://www.paddlepaddle.org.cn/packages/stable/cu118/",
"cu126": "https://www.paddlepaddle.org.cn/packages/stable/cu126/",
}
def main() -> int:
parser = argparse.ArgumentParser(
description="安装 PP-OCRv6 GPU 子项目依赖",
)
parser.add_argument(
"--cuda",
choices=sorted(PADDLE_INDEXES),
required=True,
help="目标机器的 CUDA Wheel 类型;必须依据 PaddlePaddle 官方兼容表选择",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="只显示命令,不创建环境",
)
parser.add_argument(
"--allow-no-gpu",
action="store_true",
help="允许在未检测到 nvidia-smi 时创建环境(仅用于准备/CI不代表可运行",
)
args = parser.parse_args()
uv = shutil.which("uv")
if not uv:
print("[ERROR] 未找到 uv请先安装 uv。", file=sys.stderr)
return 1
nvidia_smi = shutil.which("nvidia-smi")
if not args.dry_run and not nvidia_smi and not args.allow_no_gpu:
print(
"[ERROR] 未检测到 nvidia-smi拒绝在无 NVIDIA GPU 的机器安装 CUDA 依赖。\n"
"如仅准备环境,请显式添加 --allow-no-gpu。",
file=sys.stderr,
)
return 2
project_dir = Path(__file__).resolve().parent
index_url = PADDLE_INDEXES[args.cuda]
command = [
uv,
"sync",
"--project",
str(project_dir),
"--index",
index_url,
]
print(f"目标 CUDA Wheel: {args.cuda}")
print(f"PaddlePaddle 索引: {index_url}")
print("执行命令:")
print(" " + " ".join(command))
if args.dry_run:
return 0
ready_marker = project_dir / ".gpu-ready"
if ready_marker.exists():
ready_marker.unlink()
completed = subprocess.run(command, check=False)
if completed.returncode != 0:
print(
"[ERROR] 依赖安装失败。请检查 GPU、驱动、Python 和 PaddlePaddle Wheel 兼容性。",
file=sys.stderr,
)
return completed.returncode
ready_marker.write_text(
f"cuda={args.cuda}\nindex={index_url}\n",
encoding="utf-8",
)
print("\n[OK] GPU 子项目环境已创建。下一步从仓库根目录运行:")
print(" python ocr.py verify --device gpu")
return 0
if __name__ == "__main__":
raise SystemExit(main())

0
logs/.gitkeep Normal file
View File

68
ocr.py Normal file
View File

@ -0,0 +1,68 @@
"""PP-OCRv6 unified launcher using isolated CPU/GPU uv projects."""
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
ocr_app/__init__.py Normal file
View File

140
ocr_app/cli.py Normal file
View File

@ -0,0 +1,140 @@
"""Path-first unified CLI for PP-OCRv6."""
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_bool_option(parser: argparse.ArgumentParser, name: str, default: bool, help_text: str) -> None:
parser.add_argument(
f"--{name}",
action=argparse.BooleanOptionalAction,
default=default,
help=help_text,
)
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 _add_model_options(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--model-size",
"--model",
dest="model_size",
choices=("tiny", "small", "medium"),
default="medium",
help="PP-OCRv6 检测与识别模型规格",
)
parser.add_argument("--lang", default="ch", help="结果元数据中的语言标识,默认中文/中英混合")
parser.add_argument("--text-rec-score-thresh", type=float, default=0.0, help="识别置信度阈值")
parser.add_argument("--text-det-limit-side-len", type=int, default=None, help="文本检测边长限制")
parser.add_argument("--text-det-limit-type", choices=("min", "max"), default=None, help="检测边长限制方式")
parser.add_argument("--text-det-thresh", type=float, default=None, help="文本检测像素阈值")
parser.add_argument("--text-det-box-thresh", type=float, default=None, help="文本框阈值")
parser.add_argument("--text-det-unclip-ratio", type=float, default=None, help="文本框扩张比例")
parser.add_argument("--text-recognition-batch-size", type=int, default=6, help="文本识别批大小")
parser.add_argument("--return-word-box", action="store_true", help="返回单词级坐标")
_add_bool_option(parser, "doc-orientation-classify", True, "启用文档方向分类")
_add_bool_option(parser, "doc-unwarping", False, "启用文档去畸变")
_add_bool_option(parser, "textline-orientation", True, "启用文本行方向分类")
def build_input_parser(device_override: str | None = None) -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="ocr.py",
description="使用 PP-OCRv6 自动处理图片、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="输出根目录")
parser.add_argument("--fail-fast", action="store_true", help="单文件失败后立即停止目录任务")
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="不在日志中逐行记录识别文字")
parser.add_argument("--save-visualization", action="store_true", help="保存 OCR 可视化图片")
parser.add_argument("--save-raw-result", action="store_true", help="保存 PaddleOCR 原始 JSON")
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("--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_model_options(parser)
_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 环境")
_add_model_options(parser)
_add_device_options(parser, device_override)
return parser
def normalize_argv(argv: list[str]) -> tuple[str, list[str]]:
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
stem = "verify" if command == "verify" else (args.input.stem or args.input.name or "input")
category = "verify" if command == "verify" else "input"
log_file = args.log_file or default_log_path(PROJECT_ROOT, category, stem, device=args.device)
logger = setup_run_logger(f"ppocrv6.{category}.{args.device}", log_file, verbose=args.verbose)
provider = PipelineProvider(
RuntimeConfig(
device=args.device,
threads=args.threads,
device_id=args.device_id,
lang=args.lang,
model_size=args.model_size,
use_doc_orientation_classify=args.doc_orientation_classify,
use_doc_unwarping=args.doc_unwarping,
use_textline_orientation=args.textline_orientation,
text_recognition_batch_size=args.text_recognition_batch_size,
),
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")
if args.text_recognition_batch_size < 1:
raise ValueError("--text-recognition-batch-size 必须 >= 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

310
ocr_app/commands.py Normal file
View File

@ -0,0 +1,310 @@
"""Unified suffix-based routing for PP-OCRv6 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 .result_adapter import adapt_ocr_result
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"
if path.suffix.lower() in IMAGE_EXTENSIONS:
return "image"
if path.suffix.lower() 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 _predict_kwargs(args) -> dict[str, Any]:
return {
key: value
for key, value in {
"text_det_limit_side_len": args.text_det_limit_side_len,
"text_det_limit_type": args.text_det_limit_type,
"text_det_thresh": args.text_det_thresh,
"text_det_box_thresh": args.text_det_box_thresh,
"text_det_unclip_ratio": args.text_det_unclip_ratio,
"text_rec_score_thresh": args.text_rec_score_thresh,
"return_word_box": args.return_word_box,
}.items()
if value is not None
}
def process_image_file(
path: Path,
*,
args,
provider: PipelineProvider,
logger: logging.Logger,
project_root: Path,
run_warmup: bool,
batch_root: Path | None,
) -> FileProcessResult:
del project_root
file_started = time.perf_counter()
logger.info("FILE_ROUTED path=%s kind=image", path)
pipeline = provider.get()
predict_kwargs = _predict_kwargs(args)
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()
results = pipeline.predict(str(path), **predict_kwargs)
provider.synchronize()
elapsed = time.perf_counter() - started
inference_times.append(elapsed)
if not results:
raise RuntimeError("PP-OCRv6 未返回图片结果")
result = results[0]
logger.info("INFERENCE_COMPLETED path=%s round=%d/%d seconds=%.3f", path, index + 1, args.rounds, elapsed)
assert result is not None
adapt_started = time.perf_counter()
normalized = adapt_ocr_result(
result,
input_path=path,
source_type="image_ocr",
language=provider.config.lang,
detection_model=provider.config.text_detection_model_name,
recognition_model=provider.config.text_recognition_model_name,
model_size=provider.config.model_size,
)
adapt_seconds = time.perf_counter() - adapt_started
summary = normalized["summary"]
benchmark = {
"timestamp": datetime.now().astimezone().isoformat(),
**provider.metadata(),
"image_path": str(path),
"image": normalized["image"],
"ocr_summary": 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),
},
"result_adapt_seconds": round(adapt_seconds, 3),
"gpu_memory": provider.gpu_memory(),
"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,
normalized,
output_dir,
input_path=path,
benchmark=benchmark,
save_raw_result=args.save_raw_result,
save_visualization=args.save_visualization,
)
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 = args.benchmark_json.expanduser().resolve()
atomic_write_json(explicit, benchmark)
output_paths["explicit_benchmark"] = str(explicit)
logger.info(
"IMAGE_COMPLETED path=%s detected_lines=%d non_empty_lines=%d mean_score=%s inference_mean_seconds=%.3f file_total_seconds=%.3f output=%s",
path,
summary["detected_lines"],
summary["non_empty_lines"],
summary["mean_score"],
statistics.fmean(inference_times),
total_seconds,
output_dir,
)
if not args.no_result:
for line in normalized["lines"]:
logger.info("OCR_LINE path=%s index=%d score=%s text=%s", path, line["index"], line["score"], line["text"].replace("\n", "\\n"))
return FileProcessResult(path, "image", "completed", total_seconds, {**summary, **output_paths})
def process_pdf_file(path: Path, *, args, provider: PipelineProvider, logger: logging.Logger, batch_root: Path | None) -> FileProcessResult:
started = time.perf_counter()
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()
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", path, preflight["page_count"], len(preflight["selected_pages"]))
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=_predict_kwargs(args),
logger=logger,
)
total = time.perf_counter() - started
logger.info("PDF_COMPLETED path=%s status=%s text_pages=%d ocr_pages=%d failed_pages=%s total_seconds=%.3f", path, summary["status"], summary["text_pages"], summary["ocr_pages"], summary["failed_pages"], total)
return FileProcessResult(path, "pdf", summary["status"], total, summary, 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)
raise ValueError(f"不支持的文件类型: {path.suffix or '<无后缀>'}")
def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
program_started = time.perf_counter()
input_path = args.input.expanduser().resolve()
if detect_input_kind(input_path) != "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", 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 total_seconds=%.3f", result.path, result.kind, result.status, 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", input_path)
return 1
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:
return 130
except Exception as exc:
failures.append({"path": str(path), "error": f"{type(exc).__name__}: {exc}"})
logger.exception("FILE_FAILED path=%s", path)
if args.fail_fast:
break
program_total = time.perf_counter() - program_started
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"
atomic_write_json(
manifest_path,
{
"input_directory": str(input_path),
"recursive": args.recursive,
"device": provider.resolved_device,
"model_config": provider.model_config(),
"discovered_files": len(files),
"completed_files": len(results),
"failed_files": len(failures),
"program_total_seconds": round(program_total, 3),
"results": [{"path": str(item.path), "kind": item.kind, "status": item.status, "seconds": round(item.seconds, 3), "outputs": item.details} for item in results],
"failures": failures,
},
)
logger.info("DIRECTORY_SUMMARY discovered=%d completed=%d failed=%d manifest=%s", len(files), len(results), len(failures), manifest_path)
return 0 if not failures else 3
def run_verify(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
del args, project_root
started = time.perf_counter()
try:
provider.prepare()
paddle = provider._paddle
if provider.config.device == "gpu":
result = paddle.matmul(paddle.ones([1024, 1024]), paddle.ones([1024, 1024]))
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

74
ocr_app/logging_utils.py Normal file
View File

@ -0,0 +1,74 @@
"""Shared UTF-8 logging helpers for OCR scripts."""
from __future__ import annotations
import logging
import re
import sys
from datetime import datetime
from pathlib import Path
def safe_log_stem(value: str) -> str:
cleaned = re.sub(r"[^\w.-]+", "_", value, flags=re.UNICODE).strip("._")
return cleaned or "ocr"
def default_log_path(
project_root: Path,
category: str,
stem: str,
*,
device: str | None = None,
) -> Path:
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
suffix = f"-{safe_log_stem(device)}" if device else ""
filename = f"{safe_log_stem(stem)}{suffix}-{timestamp}.log"
return project_root / "logs" / safe_log_stem(category) / filename
def setup_run_logger(
name: str,
log_file: Path,
*,
verbose: bool = False,
console: bool = True,
) -> logging.Logger:
"""Create an isolated logger that writes UTF-8 text and optional console output."""
log_file = log_file.expanduser().resolve()
log_file.parent.mkdir(parents=True, exist_ok=True)
logger = logging.getLogger(name)
logger.setLevel(logging.DEBUG if verbose else logging.INFO)
logger.propagate = False
for handler in logger.handlers[:]:
handler.close()
logger.removeHandler(handler)
formatter = logging.Formatter(
fmt="%(asctime)s | %(levelname)-8s | pid=%(process)d | %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
file_handler = logging.FileHandler(log_file, encoding="utf-8")
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
if console:
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.DEBUG if verbose else logging.INFO)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)
logger.info("LOG_INITIALIZED file=%s", log_file)
return logger
def close_logger(logger: logging.Logger) -> None:
for handler in logger.handlers[:]:
try:
handler.flush()
handler.close()
finally:
logger.removeHandler(handler)

115
ocr_app/output.py Normal file
View File

@ -0,0 +1,115 @@
"""Output helpers for normalized PP-OCRv6 results."""
from __future__ import annotations
import json
import re
from pathlib import Path
from typing import Any
from PIL import Image
from .result_adapter import raw_result_json, result_markdown, result_plain_text
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:
base /= image_path.parent.resolve().relative_to(batch_root.resolve())
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:
base /= pdf_path.parent.resolve().relative_to(batch_root.resolve())
return base
def _save_visualization(result: Any, path: Path) -> bool:
try:
images = result.img
except Exception:
return False
if isinstance(images, dict):
image = images.get("ocr_res_img")
if image is None:
image = next(iter(images.values()), None)
else:
image = images
if image is None:
return False
if not isinstance(image, Image.Image):
try:
image = Image.fromarray(image)
except Exception:
return False
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(f".{path.name}.tmp")
image.convert("RGB").save(temporary, format="JPEG", quality=92)
temporary.replace(path)
return True
def save_image_ocr_outputs(
result: Any,
normalized: dict[str, Any],
output_dir: Path,
*,
input_path: Path,
benchmark: dict[str, Any],
save_raw_result: bool,
save_visualization: bool,
) -> dict[str, str]:
output_dir.mkdir(parents=True, exist_ok=True)
paths: dict[str, str] = {
"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"]), result_markdown(normalized, title=input_path.name))
text = result_plain_text(normalized)
atomic_write_text(Path(paths["text"]), text.rstrip() + ("\n" if text else ""))
atomic_write_json(Path(paths["json"]), normalized)
atomic_write_json(Path(paths["benchmark"]), benchmark)
if save_raw_result:
raw_path = output_dir / "raw-result.json"
atomic_write_json(raw_path, raw_result_json(result))
paths["raw_json"] = str(raw_path)
if save_visualization:
visualization = output_dir / "visualization.jpg"
if _save_visualization(result, visualization):
paths["visualization"] = str(visualization)
return paths

397
ocr_app/pdf.py Normal file
View File

@ -0,0 +1,397 @@
"""Hybrid PDF processing with PP-OCRv6 fallback for scanned pages."""
from __future__ import annotations
import hashlib
import json
import logging
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 .output import atomic_write_json, atomic_write_text, safe_stem
from .pdf_text import TextLayerPolicy, extract_page_text
from .result_adapter import adapt_ocr_result, result_markdown
MANIFEST_VERSION = 1
PAGE_SPEC_PATTERN = re.compile(r"^(\d+)(?:-(\d*)?)?$")
PDF_MODES = {"hybrid", "text", "ocr"}
def now_iso() -> str:
return datetime.now().astimezone().isoformat()
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 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 _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():
source = record.get("source_type", "unknown")
text = markdown_path.read_text(encoding="utf-8").strip()
markdown_parts.append(f"\n\n---\n\n## Page {page_number} ({source})\n\n{text}")
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,
model_config: dict[str, Any],
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 or manifest.get("pipeline") != "PP-OCRv6":
raise ValueError("manifest 与当前 PP-OCRv6 项目不兼容,请使用 --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")
if manifest.get("model_config") != model_config:
raise ValueError("PP-OCRv6 模型配置与原任务不一致,请使用原参数或 --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,
"pipeline": "PP-OCRv6",
"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,
"model_config": model_config,
"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"
document = pdfium.PdfDocument(str(pdf_path), password=password)
try:
page_count = len(document)
selected_indexes = parse_page_spec(pages, page_count)
manifest = _prepare_manifest(
pdf_path=pdf_path,
document_dir=document_dir,
page_count=page_count,
selected_pages=selected_indexes,
dpi=dpi,
mode=mode,
policy=text_policy,
resume=resume,
overwrite=overwrite,
model_config=provider.model_config(),
run_metadata={"device": provider.resolved_device},
)
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)]
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 = 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"
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"
markdown_path, json_path = _page_paths(document_dir, page_number)
if source_type == "text":
markdown_text = extracted_text.rstrip() + ("\n" if extracted_text else "")
payload = {
"schema_version": 1,
"source_type": "text",
"input_path": str(pdf_path),
"page_index": page_index,
"page_number": page_number,
"page_count": page_count,
"text": extracted_text,
"text_layer": assessment_dict,
"page_size_points": {"width": width_points, "height": height_points},
}
detected_lines = sum(bool(line.strip()) for line in extracted_text.splitlines())
mean_score = None
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("PP-OCRv6 未返回 PDF 页面结果")
payload = adapt_ocr_result(
results[0],
input_path=pdf_path,
source_type="pdf_ocr",
language=provider.config.lang,
detection_model=provider.config.text_detection_model_name,
recognition_model=provider.config.text_recognition_model_name,
model_size=provider.config.model_size,
page_index=page_index,
page_number=page_number,
)
payload["page_count"] = page_count
payload["ocr_reason"] = assessment.reason if mode == "hybrid" else "forced_ocr_mode"
payload["text_layer"] = assessment_dict
payload["render_dpi"] = dpi
markdown_text = result_markdown(payload)
detected_lines = payload["summary"]["detected_lines"]
mean_score = payload["summary"]["mean_score"]
export_started = time.perf_counter()
atomic_write_text(markdown_path, markdown_text)
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),
"detected_lines": detected_lines,
"mean_score": mean_score,
"completed_at": now_iso(),
}
except KeyboardInterrupt:
manifest["status"] = "interrupted"
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
raise
except Exception as exc:
manifest["pages"][str(page_number)] = {
"status": "failed",
"page_number": page_number,
"source_type": source_type,
"text_layer": assessment_dict,
"total_seconds": round(time.perf_counter() - started, 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()
manifest["run_metadata"] = provider.metadata()
manifest["updated_at"] = now_iso()
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
logger.info("PAGE_FINISHED page=%d source=%s progress=%d/%d", page_number, source_type, 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)
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,
"timing": {
"model_init_seconds": round(provider.model_init_seconds, 3),
"task_total_seconds": round(time.perf_counter() - task_started, 3),
},
}
manifest["updated_at"] = now_iso()
atomic_write_json(manifest_path, manifest)
rebuild_combined_outputs(document_dir, manifest)
return {"document_dir": str(document_dir), "manifest_path": str(manifest_path), "status": manifest["status"], **manifest["summary"]}
finally:
document.close()

127
ocr_app/pdf_text.py Normal file
View File

@ -0,0 +1,127 @@
"""PDF text-layer extraction and quality assessment for hybrid OCR."""
from __future__ import annotations
import re
import unicodedata
from dataclasses import asdict, dataclass
from typing import Any
@dataclass
class TextLayerPolicy:
min_chars: int = 50
min_printable_ratio: float = 0.85
min_content_ratio: float = 0.60
max_replacement_ratio: float = 0.02
min_chars_per_megapixel: float = 25.0
@dataclass
class TextLayerAssessment:
usable: bool
reason: str
raw_chars: int
non_whitespace_chars: int
printable_ratio: float
content_ratio: float
replacement_ratio: float
chars_per_megapixel: float
def to_dict(self) -> dict[str, Any]:
return asdict(self)
def normalize_text(text: str) -> str:
text = text.replace("\x00", "")
text = text.replace("\r\n", "\n").replace("\r", "\n")
lines = [re.sub(r"[ \t]+", " ", line).strip() for line in text.splitlines()]
compact_lines: list[str] = []
previous_blank = False
for line in lines:
if line:
compact_lines.append(line)
previous_blank = False
elif compact_lines and not previous_blank:
compact_lines.append("")
previous_blank = True
return "\n".join(compact_lines).strip()
def _is_content_character(character: str) -> bool:
if character.isalnum():
return True
code = ord(character)
return (
0x3400 <= code <= 0x4DBF
or 0x4E00 <= code <= 0x9FFF
or 0xF900 <= code <= 0xFAFF
or 0x3040 <= code <= 0x30FF
or 0xAC00 <= code <= 0xD7AF
)
def assess_text_layer(
text: str,
*,
width_points: float,
height_points: float,
policy: TextLayerPolicy,
) -> TextLayerAssessment:
normalized = normalize_text(text)
compact = [character for character in normalized if not character.isspace()]
count = len(compact)
if count == 0:
return TextLayerAssessment(False, "empty_text_layer", len(text), 0, 0.0, 0.0, 0.0, 0.0)
printable = sum(character.isprintable() and unicodedata.category(character) != "Cc" for character in compact)
content = sum(_is_content_character(character) for character in compact)
replacements = sum(character in {"\ufffd", "<EFBFBD>"} for character in compact)
page_megapixels = max((width_points * height_points) / 1_000_000.0, 0.01)
printable_ratio = printable / count
content_ratio = content / count
replacement_ratio = replacements / count
density = count / page_megapixels
checks = (
(count >= policy.min_chars, "too_few_characters"),
(printable_ratio >= policy.min_printable_ratio, "low_printable_ratio"),
(content_ratio >= policy.min_content_ratio, "low_content_ratio"),
(replacement_ratio <= policy.max_replacement_ratio, "high_replacement_ratio"),
(density >= policy.min_chars_per_megapixel, "low_text_density"),
)
reason = "usable_text_layer"
usable = True
for passed, failure_reason in checks:
if not passed:
usable = False
reason = failure_reason
break
return TextLayerAssessment(
usable=usable,
reason=reason,
raw_chars=len(text),
non_whitespace_chars=count,
printable_ratio=round(printable_ratio, 4),
content_ratio=round(content_ratio, 4),
replacement_ratio=round(replacement_ratio, 4),
chars_per_megapixel=round(density, 3),
)
def extract_page_text(page: Any, policy: TextLayerPolicy) -> tuple[str, TextLayerAssessment]:
text_page = page.get_textpage()
try:
raw_text = text_page.get_text_bounded()
finally:
text_page.close()
width, height = page.get_size()
normalized = normalize_text(raw_text)
assessment = assess_text_layer(
normalized,
width_points=width,
height_points=height,
policy=policy,
)
return normalized, assessment

196
ocr_app/result_adapter.py Normal file
View File

@ -0,0 +1,196 @@
"""Normalize PP-OCRv6 results into a stable project-owned schema."""
from __future__ import annotations
import math
from pathlib import Path
from typing import Any
from PIL import Image
SCHEMA_VERSION = 1
OCR_VERSION = "PP-OCRv6"
DEFAULT_MODEL_SIZE = "medium"
MODEL_VARIANTS = {
"tiny": ("PP-OCRv6_tiny_det", "PP-OCRv6_tiny_rec"),
"small": ("PP-OCRv6_small_det", "PP-OCRv6_small_rec"),
"medium": ("PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),
}
DEFAULT_DETECTION_MODEL, DEFAULT_RECOGNITION_MODEL = MODEL_VARIANTS[DEFAULT_MODEL_SIZE]
def model_names_for_size(model_size: str) -> tuple[str, str]:
try:
return MODEL_VARIANTS[model_size]
except KeyError as exc:
supported = ", ".join(MODEL_VARIANTS)
raise ValueError(f"不支持的 PP-OCRv6 模型规格: {model_size};可选: {supported}") from exc
def _payload(result: Any) -> dict[str, Any]:
if isinstance(result, dict):
return result
try:
return dict(result)
except (TypeError, ValueError) as exc:
raise TypeError(f"不支持的 PP-OCRv6 结果类型: {type(result).__name__}") from exc
def _python_value(value: Any) -> Any:
if hasattr(value, "tolist"):
value = value.tolist()
if isinstance(value, dict):
return {str(key): _python_value(item) for key, item in value.items()}
if isinstance(value, (list, tuple)):
return [_python_value(item) for item in value]
if hasattr(value, "item"):
try:
return value.item()
except (TypeError, ValueError):
pass
return value
def _polygon(value: Any) -> list[list[float | int]]:
points = _python_value(value) or []
normalized: list[list[float | int]] = []
for point in points:
if isinstance(point, (list, tuple)) and len(point) >= 2:
normalized.append([point[0], point[1]])
return normalized
def _box(value: Any, polygon: list[list[float | int]]) -> list[float | int]:
box = _python_value(value)
if isinstance(box, (list, tuple)) and len(box) >= 4:
return [box[0], box[1], box[2], box[3]]
if not polygon:
return []
xs = [point[0] for point in polygon]
ys = [point[1] for point in polygon]
return [min(xs), min(ys), max(xs), max(ys)]
def _image_size(result: dict[str, Any], input_path: Path | None) -> tuple[int | None, int | None]:
image = result.get("doc_preprocessor_res", {}).get("output_img")
shape = getattr(image, "shape", None)
if shape is not None and len(shape) >= 2:
return int(shape[1]), int(shape[0])
if input_path is not None and input_path.is_file():
try:
with Image.open(input_path) as opened:
return opened.size
except OSError:
pass
return None, None
def adapt_ocr_result(
result: Any,
*,
input_path: Path | str | None,
source_type: str,
language: str,
detection_model: str = DEFAULT_DETECTION_MODEL,
recognition_model: str = DEFAULT_RECOGNITION_MODEL,
model_size: str | None = None,
page_index: int | None = None,
page_number: int | None = None,
) -> dict[str, Any]:
raw = _payload(result)
path = Path(input_path).expanduser().resolve() if input_path is not None else None
def as_list(value: Any) -> list[Any]:
if value is None:
return []
if hasattr(value, "tolist"):
value = value.tolist()
return list(value)
texts = as_list(raw.get("rec_texts"))
scores = as_list(raw.get("rec_scores"))
polygon_values = raw.get("rec_polys")
if polygon_values is None:
polygon_values = raw.get("dt_polys")
polygons = as_list(polygon_values)
boxes = as_list(raw.get("rec_boxes"))
angles = as_list(raw.get("textline_orientation_angles"))
line_count = max(len(texts), len(scores), len(polygons), len(boxes))
lines: list[dict[str, Any]] = []
for index in range(line_count):
text = str(texts[index]) if index < len(texts) else ""
try:
score = float(scores[index]) if index < len(scores) else None
except (TypeError, ValueError):
score = None
polygon = _polygon(polygons[index]) if index < len(polygons) else []
box = _box(boxes[index] if index < len(boxes) else None, polygon)
orientation = angles[index] if index < len(angles) else None
lines.append(
{
"index": index + 1,
"text": text,
"score": round(score, 6) if score is not None and math.isfinite(score) else None,
"polygon": polygon,
"box": box,
"orientation": _python_value(orientation),
}
)
valid_scores = [line["score"] for line in lines if line["score"] is not None]
width, height = _image_size(raw, path)
resolved_page_index = page_index if page_index is not None else raw.get("page_index")
resolved_model_size = model_size
if resolved_model_size is None:
for size, names in MODEL_VARIANTS.items():
if names == (detection_model, recognition_model):
resolved_model_size = size
break
payload: dict[str, Any] = {
"schema_version": SCHEMA_VERSION,
"source_type": source_type,
"input_path": str(path) if path is not None else str(raw.get("input_path") or ""),
"page_index": resolved_page_index,
"model": {
"ocr_version": OCR_VERSION,
"model_size": resolved_model_size,
"detection_model": detection_model,
"recognition_model": recognition_model,
"language": language,
},
"image": {"width": width, "height": height},
"lines": lines,
"summary": {
"detected_lines": len(lines),
"non_empty_lines": sum(bool(line["text"].strip()) for line in lines),
"mean_score": round(sum(valid_scores) / len(valid_scores), 6) if valid_scores else None,
"min_score": min(valid_scores) if valid_scores else None,
"max_score": max(valid_scores) if valid_scores else None,
},
}
if page_number is not None:
payload["page_number"] = page_number
return payload
def result_plain_text(payload: dict[str, Any]) -> str:
return "\n".join(
str(line.get("text", "")).strip()
for line in payload.get("lines", [])
if str(line.get("text", "")).strip()
)
def result_markdown(payload: dict[str, Any], *, title: str | None = None) -> str:
text = result_plain_text(payload)
if title:
return f"# {title}\n\n{text}".rstrip() + "\n"
return text.rstrip() + ("\n" if text else "")
def raw_result_json(result: Any) -> dict[str, Any]:
raw_json = getattr(result, "json", None)
if raw_json is not None:
return _python_value(raw_json)
return _python_value(_payload(result))

175
ocr_app/runtime.py Normal file
View File

@ -0,0 +1,175 @@
"""Device validation and lazy PP-OCRv6 pipeline creation."""
from __future__ import annotations
import logging
import os
import platform
import time
from dataclasses import dataclass
from typing import Any
from .result_adapter import DEFAULT_MODEL_SIZE, OCR_VERSION, model_names_for_size
@dataclass
class RuntimeConfig:
device: str
threads: int | None = None
device_id: int = 0
lang: str = "ch"
use_doc_orientation_classify: bool = True
use_doc_unwarping: bool = False
use_textline_orientation: bool = True
text_recognition_batch_size: int = 6
model_size: str = DEFAULT_MODEL_SIZE
@property
def text_detection_model_name(self) -> str:
return model_names_for_size(self.model_size)[0]
@property
def text_recognition_model_name(self) -> str:
return model_names_for_size(self.model_size)[1]
class PipelineProvider:
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:
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)
paddle.set_device("cpu")
self._device_name = platform.processor() or "CPU"
self.logger.info("CPU_CONFIGURED threads=%d total_cores=%d", threads, total_cores)
else:
if not paddle.is_compiled_with_cuda():
raise RuntimeError("当前 PaddlePaddle 不是 CUDA 构建;不会回退到 CPU")
device_count = paddle.device.cuda.device_count()
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", self.resolved_device, self._device_name)
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 ocr_version=%s model_size=%s detection_model=%s recognition_model=%s device=%s",
OCR_VERSION,
self.config.model_size,
self.config.text_detection_model_name,
self.config.text_recognition_model_name,
self.resolved_device,
)
started = time.perf_counter()
from paddleocr import PaddleOCR
self._pipeline = PaddleOCR(
text_detection_model_name=self.config.text_detection_model_name,
text_recognition_model_name=self.config.text_recognition_model_name,
device=self.resolved_device,
use_doc_orientation_classify=self.config.use_doc_orientation_classify,
use_doc_unwarping=self.config.use_doc_unwarping,
use_textline_orientation=self.config.use_textline_orientation,
text_recognition_batch_size=self.config.text_recognition_batch_size,
)
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 = {"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
for key, name in {
"allocated_mb": "memory_allocated",
"reserved_mb": "memory_reserved",
"max_allocated_mb": "max_memory_allocated",
"max_reserved_mb": "max_memory_reserved",
}.items():
function = getattr(self._paddle.device.cuda, name, None)
if function is not None:
try:
stats[key] = round(float(function(self.config.device_id)) / (1024**2), 2)
except Exception:
pass
return stats
def model_config(self) -> dict[str, Any]:
return {
"ocr_version": OCR_VERSION,
"model_size": self.config.model_size,
"language": self.config.lang,
"detection_model": self.config.text_detection_model_name,
"recognition_model": self.config.text_recognition_model_name,
"use_doc_orientation_classify": self.config.use_doc_orientation_classify,
"use_doc_unwarping": self.config.use_doc_unwarping,
"use_textline_orientation": self.config.use_textline_orientation,
"text_recognition_batch_size": self.config.text_recognition_batch_size,
}
def metadata(self) -> dict[str, Any]:
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": self._paddle.__version__ if self._paddle is not None else None,
"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),
"model_config": self.model_config(),
}

0
outputs/.gitkeep Normal file
View File

6
tests/conftest.py Normal file
View File

@ -0,0 +1,6 @@
from pathlib import Path
import sys
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))

13
tests/test_launcher.py Normal file
View File

@ -0,0 +1,13 @@
from ocr import _requested_device
def test_default_device():
assert _requested_device(["verify"]) == "cpu"
def test_requested_gpu():
assert _requested_device(["verify", "--device", "gpu"]) == "gpu"
def test_requested_gpu_equals_syntax():
assert _requested_device(["verify", "--device=gpu"]) == "gpu"

View File

@ -0,0 +1,35 @@
import pytest
from ocr_app.cli import build_input_parser
from ocr_app.result_adapter import model_names_for_size
from ocr_app.runtime import RuntimeConfig
@pytest.mark.parametrize(
("model_size", "detection_model", "recognition_model"),
[
("tiny", "PP-OCRv6_tiny_det", "PP-OCRv6_tiny_rec"),
("small", "PP-OCRv6_small_det", "PP-OCRv6_small_rec"),
("medium", "PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),
],
)
def test_model_size_mapping(model_size, detection_model, recognition_model):
assert model_names_for_size(model_size) == (detection_model, recognition_model)
config = RuntimeConfig(device="cpu", model_size=model_size)
assert config.text_detection_model_name == detection_model
assert config.text_recognition_model_name == recognition_model
def test_model_size_defaults_to_medium():
args = build_input_parser().parse_args(["sample.png"])
assert args.model_size == "medium"
def test_model_alias_selects_tiny():
args = build_input_parser().parse_args(["sample.png", "--model", "tiny"])
assert args.model_size == "tiny"
def test_invalid_model_size_is_rejected():
with pytest.raises(ValueError):
model_names_for_size("large")

View File

@ -0,0 +1,116 @@
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 FakeResult(dict):
@property
def json(self):
return dict(self)
class FakePipeline:
def predict(self, path, **kwargs):
return [
FakeResult(
input_path=path,
page_index=0,
rec_texts=["hello OCR"],
rec_scores=[0.99],
rec_polys=[[[0, 0], [10, 0], [10, 5], [0, 5]]],
rec_boxes=[[0, 0, 10, 5]],
)
]
class FakeProvider:
class Config:
device = "cpu"
lang = "ch"
model_size = "medium"
text_detection_model_name = "PP-OCRv6_medium_det"
text_recognition_model_name = "PP-OCRv6_medium_rec"
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,
save_raw_result=False,
save_visualization=False,
text_det_limit_side_len=None,
text_det_limit_type=None,
text_det_thresh=None,
text_det_box_thresh=None,
text_det_unclip_ratio=None,
text_rec_score_thresh=0.0,
return_word_box=False,
)
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 "hello OCR" in (output_dir / "result.md").read_text("utf-8")
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["source_type"] == "image_ocr"
assert data["lines"][0]["score"] == 0.99
benchmark = json.loads((output_dir / "benchmark.json").read_text("utf-8"))
assert benchmark["file_total_seconds"] >= benchmark["export_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"
assert image_output_directory(output, tmp_path / "same.png", batch_root=None, recursive=False) != image_output_directory(output, tmp_path / "same.jpg", batch_root=None, recursive=False)

13
tests/test_page_spec.py Normal file
View File

@ -0,0 +1,13 @@
import pytest
from ocr_app.pdf import parse_page_spec
def test_page_ranges():
assert parse_page_spec("1-2,4,6-", 7) == [0, 1, 3, 5, 6]
@pytest.mark.parametrize("value", ["0", "3-2", "1-a", "8"])
def test_invalid_page_ranges(value):
with pytest.raises(ValueError):
parse_page_spec(value, 7)

89
tests/test_pdf_hybrid.py Normal file
View File

@ -0,0 +1,89 @@
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 FakeResult(dict):
pass
class FakePipeline:
def predict(self, path, **kwargs):
return [
FakeResult(
input_path=path,
page_index=0,
rec_texts=["mock OCR"],
rec_scores=[0.97],
rec_polys=[[[0, 0], [20, 0], [20, 10], [0, 10]]],
rec_boxes=[[0, 0, 20, 10]],
)
]
class FakeProvider:
resolved_device = "cpu"
class Config:
lang = "ch"
model_size = "medium"
text_detection_model_name = "PP-OCRv6_medium_det"
text_recognition_model_name = "PP-OCRv6_medium_rec"
config = Config()
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 model_config(self):
return {
"ocr_version": "PP-OCRv6",
"model_size": "medium",
"language": "ch",
"detection_model": "PP-OCRv6_medium_det",
"recognition_model": "PP-OCRv6_medium_rec",
}
def metadata(self):
return {"device": "cpu", "model_init_seconds": self.model_init_seconds, "model_config": self.model_config()}
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 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"))
page_json = json.loads((Path(result["document_dir"]) / "pages" / "page-0001.json").read_text("utf-8"))
assert result["ocr_pages"] == 1
assert provider.get_calls == 1
assert manifest["pages"]["1"]["routing_reason"] == "empty_text_layer"
assert page_json["lines"][0]["text"] == "mock OCR"

39
tests/test_pdf_text.py Normal file
View File

@ -0,0 +1,39 @@
from ocr_app.pdf_text import TextLayerPolicy, assess_text_layer, normalize_text
def test_normalize_text():
assert normalize_text("a b\r\n\r\n\r\nc") == "a b\n\nc"
def test_usable_text_layer():
text = "有效电子文档内容 123 " * 20
result = assess_text_layer(
text,
width_points=595,
height_points=842,
policy=TextLayerPolicy(),
)
assert result.usable
assert result.reason == "usable_text_layer"
def test_empty_text_layer_routes_to_ocr():
result = assess_text_layer(
"",
width_points=595,
height_points=842,
policy=TextLayerPolicy(),
)
assert not result.usable
assert result.reason == "empty_text_layer"
def test_short_text_layer_routes_to_ocr():
result = assess_text_layer(
"页码 1",
width_points=595,
height_points=842,
policy=TextLayerPolicy(min_chars=50),
)
assert not result.usable
assert result.reason == "too_few_characters"

View File

@ -0,0 +1,36 @@
from pathlib import Path
from PIL import Image
from ocr_app.result_adapter import adapt_ocr_result, result_markdown, result_plain_text
def test_adapt_ppocr_result(tmp_path):
image = tmp_path / "sample.png"
Image.new("RGB", (100, 50), "white").save(image)
result = {
"page_index": 0,
"rec_texts": ["第一行", "第二行"],
"rec_scores": [0.98, 0.75],
"rec_polys": [
[[1, 2], [20, 2], [20, 10], [1, 10]],
[[2, 20], [30, 20], [30, 30], [2, 30]],
],
"rec_boxes": [[1, 2, 20, 10], [2, 20, 30, 30]],
}
payload = adapt_ocr_result(result, input_path=image, source_type="image_ocr", language="ch")
assert payload["schema_version"] == 1
assert payload["image"] == {"width": 100, "height": 50}
assert payload["summary"]["detected_lines"] == 2
assert payload["summary"]["mean_score"] == 0.865
assert payload["lines"][0]["text"] == "第一行"
assert result_plain_text(payload) == "第一行\n第二行"
assert result_markdown(payload, title="sample.png").startswith("# sample.png\n\n第一行")
def test_empty_result_is_supported(tmp_path):
payload = adapt_ocr_result({}, input_path=Path(tmp_path / "missing.png"), source_type="image_ocr", language="ch")
assert payload["lines"] == []
assert payload["summary"]["mean_score"] is None