feat: 新增 pdf ocr 能力
This commit is contained in:
parent
afb487260a
commit
b3207088b7
@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
name: pdf
|
name: pdf
|
||||||
description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全下载、元数据与页面检查、多引擎分段文本提取和质量检测、表格提取、按需页面 PNG 渲染、从文本创建 PDF、合并、拆分、旋转及最终质量校验。当用户提供 .pdf 文件或 HTTPS PDF 地址,或要求总结、读取、生成、编辑、转换或审阅 PDF 时使用。"
|
description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全下载、元数据与页面检查、多引擎分段文本提取和质量检测、扫描页本地 OCR、表格提取、按需页面 PNG 渲染、从文本创建 PDF、合并、拆分、旋转及最终质量校验。当用户提供 .pdf 文件或 HTTPS PDF 地址,或要求总结、读取、识别扫描件、生成、编辑、转换或审阅 PDF 时使用。"
|
||||||
---
|
---
|
||||||
|
|
||||||
# PDF 处理
|
# PDF 处理
|
||||||
@ -15,7 +15,7 @@ description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全
|
|||||||
- 不创建或猜测脚本清单以外的文件。
|
- 不创建或猜测脚本清单以外的文件。
|
||||||
- 每次检查脚本返回的 JSON;只有 `ok` 为 `true` 时才继续。
|
- 每次检查脚本返回的 JSON;只有 `ok` 为 `true` 时才继续。
|
||||||
- 收到 `ok: false` 时,依据 `error` 调整合法参数或向用户说明失败原因,不要把参数改传给其他脚本碰运气。
|
- 收到 `ok: false` 时,依据 `error` 调整合法参数或向用户说明失败原因,不要把参数改传给其他脚本碰运气。
|
||||||
- 阅读或总结时只使用 `pages[]` 中 `usable_for_summary: true` 的文本。`needs_ocr: false` 时不得为了“常规检查”继续渲染或调用图片识别。
|
- 阅读或总结时只使用 `pages[]` 中 `usable_for_summary: true` 的文本。`needs_ocr: false` 时不得为了“常规检查”继续 OCR、渲染或调用图片识别。
|
||||||
- 不把 PDF 密码作为脚本参数;工具调用参数可能进入运行日志。
|
- 不把 PDF 密码作为脚本参数;工具调用参数可能进入运行日志。
|
||||||
|
|
||||||
## 脚本清单
|
## 脚本清单
|
||||||
@ -25,6 +25,7 @@ description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全
|
|||||||
| `scripts/download_pdf.py` | 下载并校验远程 HTTPS PDF | `urllib`、`pypdf` |
|
| `scripts/download_pdf.py` | 下载并校验远程 HTTPS PDF | `urllib`、`pypdf` |
|
||||||
| `scripts/inspect_pdf.py` | 检查页数、加密、元数据、页面尺寸和表单数量 | `pypdf` |
|
| `scripts/inspect_pdf.py` | 检查页数、加密、元数据、页面尺寸和表单数量 | `pypdf` |
|
||||||
| `scripts/extract_text.py` | 多引擎提取、质量检测并分段返回正文 | Poppler `pdftotext`、`pdfplumber`;`pypdf` 校验 |
|
| `scripts/extract_text.py` | 多引擎提取、质量检测并分段返回正文 | Poppler `pdftotext`、`pdfplumber`;`pypdf` 校验 |
|
||||||
|
| `scripts/ocr_text.py` | 对指定扫描页执行离线 OCR 并返回可靠文字 | RapidOCR、ONNX Runtime、Poppler `pdftoppm` |
|
||||||
| `scripts/extract_tables.py` | 按页提取表格 | `pdfplumber` |
|
| `scripts/extract_tables.py` | 按页提取表格 | `pdfplumber` |
|
||||||
| `scripts/render_pdf.py` | 把指定页面渲染为 PNG | Poppler `pdftoppm` |
|
| `scripts/render_pdf.py` | 把指定页面渲染为 PNG | Poppler `pdftoppm` |
|
||||||
| `scripts/create_pdf.py` | 从 UTF-8 文本或 Markdown 创建 PDF | `reportlab`、`pypdf` |
|
| `scripts/create_pdf.py` | 从 UTF-8 文本或 Markdown 创建 PDF | `reportlab`、`pypdf` |
|
||||||
@ -38,11 +39,12 @@ description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全
|
|||||||
1. 为任务选择简短目录名,把中间文件放在 `tmp/pdfs/<任务名>/`。
|
1. 为任务选择简短目录名,把中间文件放在 `tmp/pdfs/<任务名>/`。
|
||||||
2. 远程 HTTPS 链接先调用 `download_pdf.py`;本地文件直接进入下一步。
|
2. 远程 HTTPS 链接先调用 `download_pdf.py`;本地文件直接进入下一步。
|
||||||
3. 调用 `inspect_pdf.py` 检查文件。遇到加密 PDF 时停止处理,请用户提供已解密副本;当前固定脚本不接收密码。
|
3. 调用 `inspect_pdf.py` 检查文件。遇到加密 PDF 时停止处理,请用户提供已解密副本;当前固定脚本不接收密码。
|
||||||
4. 阅读或总结时调用 `extract_text.py`。结果为 `usable_for_summary: true` 时使用可靠页文本并根据游标继续;同时为 `needs_ocr: false` 时直接回答,不调用渲染或图片识别。
|
4. 阅读或总结时调用 `extract_text.py`。结果为 `usable_for_summary: true` 时使用可靠页文本并根据游标继续;同时为 `needs_ocr: false` 时直接回答,不调用 OCR、渲染或图片识别。
|
||||||
5. 仅在文本结果为 `needs_ocr: true`、用户明确要求检查视觉版式,或任务涉及创建/修改 PDF 时调用 `render_pdf.py`。
|
5. 只有 `extract_text.py` 返回 `needs_ocr: true` 时,才对 `text_quality.suspect_pages` 调用 `ocr_text.py`。原生可靠文本优先,OCR 只补齐可疑页,不重复识别正常页。
|
||||||
6. 文本兜底时先只渲染 1 个有代表性的页面确认是否为扫描内容;确认后才按任务所需页面继续渲染和识别,不默认渲染前 5 页。
|
6. `ocr_text.py` 会在脚本内部临时渲染指定页面并交给本地 RapidOCR,完成后自动删除 PNG;普通扫描件解析不调用大模型识图,也不需要先调用 `render_pdf.py`。
|
||||||
7. 创建或修改后的最终 PDF 写入 `output/pdf/`,重新执行检查、文本提取和全部页面渲染。
|
7. 仅在用户明确要求检查视觉版式,或任务涉及创建/修改 PDF 时调用 `render_pdf.py`。
|
||||||
8. 最终产物位于临时目录之外且不再需要缓存时,调用 `cleanup_pdf_temp.py` 清理本次任务目录。
|
8. 创建或修改后的最终 PDF 写入 `output/pdf/`,重新执行检查、文本提取和全部页面渲染。
|
||||||
|
9. 最终产物位于临时目录之外且不再需要缓存时,调用 `cleanup_pdf_temp.py` 清理本次任务目录。
|
||||||
|
|
||||||
## 下载远程 PDF
|
## 下载远程 PDF
|
||||||
|
|
||||||
@ -95,10 +97,35 @@ description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全
|
|||||||
|
|
||||||
- `usable_for_summary: true`:只使用 `pages[]` 中同样标为 `usable_for_summary: true` 的 `text`;可疑页的文本会被置空。如果 `has_more: true`,始终传回 `next_page` 和 `next_offset`。仅当 `next_offset` 大于 `0` 时,把非空的 `next_engine` 传给 `--engine` 以固定同页字符游标;这种调用只续读当前页。当前页完成后返回的 `next_offset` 为 `0`,此时不要传 `--engine`,让下一页重新使用 `auto`。保留首次调用的 `--end-page`(如果指定)及其他选项,直至 `has_more: false`。
|
- `usable_for_summary: true`:只使用 `pages[]` 中同样标为 `usable_for_summary: true` 的 `text`;可疑页的文本会被置空。如果 `has_more: true`,始终传回 `next_page` 和 `next_offset`。仅当 `next_offset` 大于 `0` 时,把非空的 `next_engine` 传给 `--engine` 以固定同页字符游标;这种调用只续读当前页。当前页完成后返回的 `next_offset` 为 `0`,此时不要传 `--engine`,让下一页重新使用 `auto`。保留首次调用的 `--end-page`(如果指定)及其他选项,直至 `has_more: false`。
|
||||||
- `usable_for_summary: false`:本批次没有可靠文本,不要使用返回内容。查看 `engine_attempts`、`text_quality.reasons`、`text_quality.suspect_pages` 和 `needs_ocr`;若 `has_more: true`,仍按跨页游标继续检查后续批次,避免漏掉后续可搜索文本。
|
- `usable_for_summary: false`:本批次没有可靠文本,不要使用返回内容。查看 `engine_attempts`、`text_quality.reasons`、`text_quality.suspect_pages` 和 `needs_ocr`;若 `has_more: true`,仍按跨页游标继续检查后续批次,避免漏掉后续可搜索文本。
|
||||||
- `needs_ocr: true`:一个或多个页面未得到可靠文本。优先选择 `text_quality.suspect_pages` 中的首个页码,只渲染该页诊断;确认确为扫描内容后才调用图片识别。
|
- `needs_ocr: true`:一个或多个页面未得到可靠文本。把 `text_quality.suspect_pages` 中实际需要阅读的页码传给 `ocr_text.py`;不要先调用 `render_pdf.py`,也不要把临时图片交给大模型。
|
||||||
|
|
||||||
`complete_text_coverage: true` 表示本批次所有页面均有可靠文本。`text_quality` 按页检测空白或过少文本、页面实际可见图像覆盖过大但文字不足、`(cid:...)`、Unicode 替换字符、异常控制字符及外观像汉字的部首字符;`pages[].extractor` 表示该页最终采用的引擎。`status: mixed` 表示同一批次同时包含可靠页和可疑页:可先使用可靠页文本,同时只核验 `suspect_pages`。不要只根据“肉眼看起来能读”判定提取结果可靠。
|
`complete_text_coverage: true` 表示本批次所有页面均有可靠文本。`text_quality` 按页检测空白或过少文本、页面实际可见图像覆盖过大但文字不足、`(cid:...)`、Unicode 替换字符、异常控制字符及外观像汉字的部首字符;`pages[].extractor` 表示该页最终采用的引擎。`status: mixed` 表示同一批次同时包含可靠页和可疑页:可先使用可靠页文本,同时只核验 `suspect_pages`。不要只根据“肉眼看起来能读”判定提取结果可靠。
|
||||||
|
|
||||||
|
## 本地 OCR 扫描页
|
||||||
|
|
||||||
|
仅当 `extract_text.py` 返回 `needs_ocr: true` 时调用 `scripts/ocr_text.py`。`--pages` 必须明确指定 `text_quality.suspect_pages` 中要读取的页,单次最多 4 页:
|
||||||
|
|
||||||
|
```text
|
||||||
|
--input 'tmp/pdfs/<任务名>/source.pdf' --pages '2,5-6'
|
||||||
|
```
|
||||||
|
|
||||||
|
默认以 260 DPI 临时渲染,并使用镜像中预置的 RapidOCR 与 ONNX Runtime 在本地识别。脚本不会联网下载模型,不会保留渲染图片,也不会调用大模型视觉能力。可选参数:
|
||||||
|
|
||||||
|
- `--dpi <150-400>`:文字过小或识别质量不足时适度提高,默认 `260`。
|
||||||
|
- `--max-chars <字符数>`:默认 `24000`,最大 `60000`。
|
||||||
|
- `--timeout <秒>`:每页 Poppler 渲染超时,默认 `180`。
|
||||||
|
- `--start-offset <字符偏移>`:续读被字符上限截断的单页;使用时 `--pages` 只能包含该页。
|
||||||
|
|
||||||
|
只使用 `pages[]` 中 `usable_for_summary: true` 的 `text`。`status: empty`、`sparse` 或 `low_confidence` 的页面文本会被置空,并通过 `needs_review: true` 提醒人工检查。
|
||||||
|
|
||||||
|
如果 `has_more: true`:
|
||||||
|
|
||||||
|
- `next_offset > 0`:用 `--pages <next_page> --start-offset <next_offset>` 续读同一页。
|
||||||
|
- `next_offset = 0`:用返回的 `remaining_pages` 继续下一批。
|
||||||
|
- 同页续读完成后,再处理先前返回的其他 `remaining_pages`。
|
||||||
|
|
||||||
|
OCR 结果中的 `mean_confidence`、`line_count`、`render_seconds` 和 `ocr_seconds` 仅用于判断质量与性能。原生提取成功的页面始终采用 `extract_text.py` 结果,不用 OCR 覆盖。
|
||||||
|
|
||||||
## 提取表格
|
## 提取表格
|
||||||
|
|
||||||
调用 `scripts/extract_tables.py`:
|
调用 `scripts/extract_tables.py`:
|
||||||
@ -113,11 +140,10 @@ description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全
|
|||||||
|
|
||||||
只有满足以下任一条件时才调用 `scripts/render_pdf.py`:
|
只有满足以下任一条件时才调用 `scripts/render_pdf.py`:
|
||||||
|
|
||||||
- `extract_text.py` 返回 `needs_ocr: true`;
|
|
||||||
- 用户明确要求审阅版式、图表、印章、公式或页面外观;
|
- 用户明确要求审阅版式、图表、印章、公式或页面外观;
|
||||||
- 创建或修改 PDF 后进行最终视觉检查。
|
- 创建或修改 PDF 后进行最终视觉检查。
|
||||||
|
|
||||||
不要因为输入是 PDF、需要总结或需要检查首页就自动渲染。调用脚本时不要直接执行 `pdftoppm`:
|
不要因为输入是 PDF、需要总结、需要 OCR 或需要检查首页就自动调用本脚本;OCR 的临时渲染由 `ocr_text.py` 内部完成。调用脚本时不要直接执行 `pdftoppm`:
|
||||||
|
|
||||||
```text
|
```text
|
||||||
--input 'tmp/pdfs/<任务名>/source.pdf' --output-dir 'tmp/pdfs/<任务名>/rendered' --start-page 1
|
--input 'tmp/pdfs/<任务名>/source.pdf' --output-dir 'tmp/pdfs/<任务名>/rendered' --start-page 1
|
||||||
@ -125,7 +151,7 @@ description: "处理本地 PDF 文件或远程 HTTPS PDF 链接,包括安全
|
|||||||
|
|
||||||
默认 150 DPI、单次最多 10 页。可使用 `--end-page`、`--max-pages`、`--dpi`、`--timeout` 和 `--overwrite`。若 `has_more: true`,使用 `next_page` 继续,并保留首次调用的 `--end-page`(如果指定)、输出目录及其他渲染选项。脚本返回标准化的 `page-0001.png` 文件路径。
|
默认 150 DPI、单次最多 10 页。可使用 `--end-page`、`--max-pages`、`--dpi`、`--timeout` 和 `--overwrite`。若 `has_more: true`,使用 `next_page` 继续,并保留首次调用的 `--end-page`(如果指定)、输出目录及其他渲染选项。脚本返回标准化的 `page-0001.png` 文件路径。
|
||||||
|
|
||||||
文本兜底首次只传一个页码,例如同时设置 `--start-page 2 --end-page 2 --max-pages 1`。对文字较小或图表密集的页面提高 DPI。使用可用的图像查看或识别工具检查返回的 PNG,不要尝试把图片路径交给下载脚本。
|
对文字较小或图表密集的页面提高 DPI。使用可用的图像查看工具检查返回的 PNG,不要尝试把图片路径交给下载脚本。
|
||||||
|
|
||||||
## 创建 PDF
|
## 创建 PDF
|
||||||
|
|
||||||
@ -185,6 +211,7 @@ rotate --input 'source.pdf' --output 'output/pdf/rotated.pdf' --pages '1,3-5' --
|
|||||||
|
|
||||||
- 不覆盖用户提供的源文件。
|
- 不覆盖用户提供的源文件。
|
||||||
- 创建或修改后重新检查页数、页面尺寸、加密状态和文本可读性。
|
- 创建或修改后重新检查页数、页面尺寸、加密状态和文本可读性。
|
||||||
|
- 扫描件先做原生文字检测,再只 OCR 可疑页;不得把低置信度 OCR 文本当作可靠正文。
|
||||||
- 逐页确认没有裁切、重叠、溢出、乱码、黑方块、错误分页或异常空白页。
|
- 逐页确认没有裁切、重叠、溢出、乱码、黑方块、错误分页或异常空白页。
|
||||||
- 检查标题层级、段落间距、页边距、表格、图表、图片、页码及章节衔接。
|
- 检查标题层级、段落间距、页边距、表格、图表、图片、页码及章节衔接。
|
||||||
- 引用和参考文献必须可读,不得残留工具令牌、占位符或临时路径。
|
- 引用和参考文献必须可读,不得残留工具令牌、占位符或临时路径。
|
||||||
|
|||||||
@ -1,4 +1,4 @@
|
|||||||
interface:
|
interface:
|
||||||
display_name: "PDF 处理"
|
display_name: "PDF 处理"
|
||||||
short_description: "读取、创建和审阅本地或远程 PDF,文本优先并按需检查版式"
|
short_description: "读取、创建和审阅本地或远程 PDF,按需执行本地 OCR"
|
||||||
default_prompt: "使用 $pdf 下载或读取这个 PDF,优先提取可靠文本并按我的要求处理内容。"
|
default_prompt: "使用 $pdf 下载或读取这个 PDF,优先提取可靠文本,并只对扫描页执行本地 OCR。"
|
||||||
|
|||||||
456
skills/pdf/scripts/ocr_text.py
Normal file
456
skills/pdf/scripts/ocr_text.py
Normal file
@ -0,0 +1,456 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import contextlib
|
||||||
|
import importlib.metadata
|
||||||
|
import io
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import shutil
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from _pdf_common import (
|
||||||
|
SkillArgumentParser,
|
||||||
|
input_pdf,
|
||||||
|
parse_page_spec,
|
||||||
|
run_cli,
|
||||||
|
)
|
||||||
|
|
||||||
|
DEFAULT_DPI = 260
|
||||||
|
DEFAULT_MAX_CHARS = 24000
|
||||||
|
DEFAULT_RENDER_TIMEOUT_SECONDS = 180
|
||||||
|
MAX_PAGES_PER_CALL = 4
|
||||||
|
MAX_PIXELS_PER_PAGE = 20_000_000
|
||||||
|
MIN_MEAN_CONFIDENCE = 0.60
|
||||||
|
MIN_MEANINGFUL_CHARS = 5
|
||||||
|
WHITESPACE_PATTERN = re.compile(r"[ \t]+")
|
||||||
|
|
||||||
|
|
||||||
|
for variable, value in (
|
||||||
|
("OMP_NUM_THREADS", "2"),
|
||||||
|
("OPENBLAS_NUM_THREADS", "1"),
|
||||||
|
("MKL_NUM_THREADS", "1"),
|
||||||
|
("NUMEXPR_NUM_THREADS", "1"),
|
||||||
|
):
|
||||||
|
os.environ.setdefault(variable, value)
|
||||||
|
|
||||||
|
for logger_name in ("rapidocr", "RapidOCR", "onnxruntime"):
|
||||||
|
logging.getLogger(logger_name).setLevel(logging.ERROR)
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_args(argv: list[str]):
|
||||||
|
parser = SkillArgumentParser(description="本地 OCR 提取指定 PDF 页面文字")
|
||||||
|
parser.add_argument("--input", required=True, help="本地 PDF 路径")
|
||||||
|
parser.add_argument(
|
||||||
|
"--pages",
|
||||||
|
required=True,
|
||||||
|
help="要识别的页码,例如 2 或 2,5-6;单次最多 4 页",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--start-offset",
|
||||||
|
type=int,
|
||||||
|
default=0,
|
||||||
|
help="续读单页 OCR 文字时的字符偏移量",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--max-chars",
|
||||||
|
type=int,
|
||||||
|
default=DEFAULT_MAX_CHARS,
|
||||||
|
help=f"单次最多返回文本字符数,默认 {DEFAULT_MAX_CHARS}",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--dpi",
|
||||||
|
type=int,
|
||||||
|
default=DEFAULT_DPI,
|
||||||
|
help=f"OCR 渲染分辨率,默认 {DEFAULT_DPI} DPI",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--timeout",
|
||||||
|
type=int,
|
||||||
|
default=DEFAULT_RENDER_TIMEOUT_SECONDS,
|
||||||
|
help=(
|
||||||
|
"每页 Poppler 渲染超时秒数,"
|
||||||
|
f"默认 {DEFAULT_RENDER_TIMEOUT_SECONDS}"
|
||||||
|
),
|
||||||
|
)
|
||||||
|
args = parser.parse_args(argv)
|
||||||
|
if args.start_offset < 0:
|
||||||
|
raise ValueError("start-offset 不能小于 0")
|
||||||
|
if args.max_chars < 1 or args.max_chars > 60000:
|
||||||
|
raise ValueError("max-chars 必须在 1 到 60000 之间")
|
||||||
|
if args.dpi < 150 or args.dpi > 400:
|
||||||
|
raise ValueError("dpi 必须在 150 到 400 之间")
|
||||||
|
if args.timeout < 1 or args.timeout > 600:
|
||||||
|
raise ValueError("timeout 必须在 1 到 600 之间")
|
||||||
|
return args
|
||||||
|
|
||||||
|
|
||||||
|
def _pdf_pages(path: Path) -> tuple[int, dict[int, tuple[float, float]]]:
|
||||||
|
from pypdf import PdfReader
|
||||||
|
|
||||||
|
page_sizes: dict[int, tuple[float, float]] = {}
|
||||||
|
with path.open("rb") as stream:
|
||||||
|
reader = PdfReader(stream, strict=False)
|
||||||
|
if reader.is_encrypted:
|
||||||
|
raise ValueError("PDF 已加密,无法执行 OCR")
|
||||||
|
page_count = len(reader.pages)
|
||||||
|
for page_number, page in enumerate(reader.pages, start=1):
|
||||||
|
page_sizes[page_number] = (
|
||||||
|
abs(float(page.cropbox.width)),
|
||||||
|
abs(float(page.cropbox.height)),
|
||||||
|
)
|
||||||
|
return page_count, page_sizes
|
||||||
|
|
||||||
|
|
||||||
|
def _find_pdftoppm() -> Path:
|
||||||
|
executable = shutil.which("pdftoppm")
|
||||||
|
if not executable:
|
||||||
|
raise RuntimeError("环境预置的 pdftoppm 不可用")
|
||||||
|
return Path(executable).resolve()
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_render_size(
|
||||||
|
page_number: int,
|
||||||
|
page_size: tuple[float, float],
|
||||||
|
dpi: int,
|
||||||
|
) -> None:
|
||||||
|
width_points, height_points = page_size
|
||||||
|
estimated_pixels = (
|
||||||
|
width_points * dpi / 72.0
|
||||||
|
* height_points * dpi / 72.0
|
||||||
|
)
|
||||||
|
if estimated_pixels > MAX_PIXELS_PER_PAGE:
|
||||||
|
raise ValueError(
|
||||||
|
f"第 {page_number} 页按 {dpi} DPI 渲染预计超过 "
|
||||||
|
f"{MAX_PIXELS_PER_PAGE} 像素,请降低 dpi"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _render_page(
|
||||||
|
executable: Path,
|
||||||
|
pdf_path: Path,
|
||||||
|
page_number: int,
|
||||||
|
page_size: tuple[float, float],
|
||||||
|
dpi: int,
|
||||||
|
timeout: int,
|
||||||
|
temp_dir: Path,
|
||||||
|
) -> tuple[Path, float]:
|
||||||
|
_validate_render_size(page_number, page_size, dpi)
|
||||||
|
prefix = temp_dir / f"page-{page_number:04d}"
|
||||||
|
output = prefix.with_suffix(".png")
|
||||||
|
command = [
|
||||||
|
str(executable),
|
||||||
|
"-f",
|
||||||
|
str(page_number),
|
||||||
|
"-l",
|
||||||
|
str(page_number),
|
||||||
|
"-singlefile",
|
||||||
|
"-png",
|
||||||
|
"-r",
|
||||||
|
str(dpi),
|
||||||
|
str(pdf_path),
|
||||||
|
str(prefix),
|
||||||
|
]
|
||||||
|
started = time.monotonic()
|
||||||
|
try:
|
||||||
|
completed = subprocess.run(
|
||||||
|
command,
|
||||||
|
text=True,
|
||||||
|
capture_output=True,
|
||||||
|
timeout=timeout,
|
||||||
|
check=False,
|
||||||
|
)
|
||||||
|
except subprocess.TimeoutExpired as exc:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"第 {page_number} 页渲染超过 {timeout} 秒"
|
||||||
|
) from exc
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
if completed.returncode != 0:
|
||||||
|
detail = (completed.stderr or completed.stdout or "").strip()[-2000:]
|
||||||
|
raise RuntimeError(
|
||||||
|
f"第 {page_number} 页渲染失败:"
|
||||||
|
f"{detail or 'pdftoppm 返回错误'}"
|
||||||
|
)
|
||||||
|
if not output.is_file() or output.stat().st_size <= 0:
|
||||||
|
raise RuntimeError(f"第 {page_number} 页没有生成有效 PNG")
|
||||||
|
return output, elapsed
|
||||||
|
|
||||||
|
|
||||||
|
def _create_ocr_engine():
|
||||||
|
try:
|
||||||
|
from rapidocr import RapidOCR
|
||||||
|
except ImportError as exc:
|
||||||
|
raise RuntimeError("环境预置的 rapidocr 模块不可用") from exc
|
||||||
|
|
||||||
|
captured_stdout = io.StringIO()
|
||||||
|
captured_stderr = io.StringIO()
|
||||||
|
with (
|
||||||
|
contextlib.redirect_stdout(captured_stdout),
|
||||||
|
contextlib.redirect_stderr(captured_stderr),
|
||||||
|
):
|
||||||
|
return RapidOCR()
|
||||||
|
|
||||||
|
|
||||||
|
def _clean_text(value: Any) -> str:
|
||||||
|
text = str(value or "").replace("\x00", "").strip()
|
||||||
|
return "\n".join(
|
||||||
|
WHITESPACE_PATTERN.sub(" ", line).strip()
|
||||||
|
for line in text.replace("\r\n", "\n").replace("\r", "\n").split("\n")
|
||||||
|
if line.strip()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _box_points(value: Any) -> list[list[float]] | None:
|
||||||
|
if value is None:
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
points = [
|
||||||
|
[round(float(point[0]), 2), round(float(point[1]), 2)]
|
||||||
|
for point in value
|
||||||
|
]
|
||||||
|
except (IndexError, TypeError, ValueError):
|
||||||
|
return None
|
||||||
|
return points if len(points) == 4 else None
|
||||||
|
|
||||||
|
|
||||||
|
def _ordered_lines(result: Any) -> list[dict[str, Any]]:
|
||||||
|
texts = list(getattr(result, "txts", None) or ())
|
||||||
|
scores = list(getattr(result, "scores", None) or ())
|
||||||
|
raw_boxes = getattr(result, "boxes", None)
|
||||||
|
boxes = list(raw_boxes) if raw_boxes is not None else []
|
||||||
|
|
||||||
|
lines: list[dict[str, Any]] = []
|
||||||
|
for index, raw_text in enumerate(texts):
|
||||||
|
text = _clean_text(raw_text)
|
||||||
|
if not text:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
confidence = float(scores[index])
|
||||||
|
except (IndexError, TypeError, ValueError):
|
||||||
|
confidence = 0.0
|
||||||
|
confidence = max(0.0, min(1.0, confidence))
|
||||||
|
box = _box_points(boxes[index] if index < len(boxes) else None)
|
||||||
|
if box:
|
||||||
|
left = min(point[0] for point in box)
|
||||||
|
top = min(point[1] for point in box)
|
||||||
|
else:
|
||||||
|
left = float(index)
|
||||||
|
top = float(index)
|
||||||
|
lines.append(
|
||||||
|
{
|
||||||
|
"text": text,
|
||||||
|
"confidence": confidence,
|
||||||
|
"box": box,
|
||||||
|
"_left": left,
|
||||||
|
"_top": top,
|
||||||
|
"_index": index,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
lines.sort(
|
||||||
|
key=lambda line: (
|
||||||
|
round(line["_top"] / 10.0),
|
||||||
|
line["_left"],
|
||||||
|
line["_index"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return lines
|
||||||
|
|
||||||
|
|
||||||
|
def _ocr_page(engine: Any, image_path: Path) -> dict[str, Any]:
|
||||||
|
captured_stdout = io.StringIO()
|
||||||
|
captured_stderr = io.StringIO()
|
||||||
|
started = time.monotonic()
|
||||||
|
with (
|
||||||
|
contextlib.redirect_stdout(captured_stdout),
|
||||||
|
contextlib.redirect_stderr(captured_stderr),
|
||||||
|
):
|
||||||
|
result = engine(str(image_path))
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
lines = _ordered_lines(result)
|
||||||
|
text = "\n".join(line["text"] for line in lines)
|
||||||
|
weighted_chars = [
|
||||||
|
max(1, sum(1 for character in line["text"] if not character.isspace()))
|
||||||
|
for line in lines
|
||||||
|
]
|
||||||
|
total_weight = sum(weighted_chars)
|
||||||
|
mean_confidence = (
|
||||||
|
sum(
|
||||||
|
line["confidence"] * weight
|
||||||
|
for line, weight in zip(lines, weighted_chars)
|
||||||
|
)
|
||||||
|
/ total_weight
|
||||||
|
if total_weight
|
||||||
|
else 0.0
|
||||||
|
)
|
||||||
|
meaningful_chars = sum(1 for character in text if character.isalnum())
|
||||||
|
low_confidence_lines = sum(
|
||||||
|
1 for line in lines if line["confidence"] < MIN_MEAN_CONFIDENCE
|
||||||
|
)
|
||||||
|
|
||||||
|
reasons: list[str] = []
|
||||||
|
if not text:
|
||||||
|
status = "empty"
|
||||||
|
reasons.append("本地 OCR 未识别到文字")
|
||||||
|
elif meaningful_chars < MIN_MEANINGFUL_CHARS:
|
||||||
|
status = "sparse"
|
||||||
|
reasons.append(
|
||||||
|
f"有效文字少于 {MIN_MEANINGFUL_CHARS} 个字符"
|
||||||
|
)
|
||||||
|
elif mean_confidence < MIN_MEAN_CONFIDENCE:
|
||||||
|
status = "low_confidence"
|
||||||
|
reasons.append(
|
||||||
|
"OCR 平均置信度低于 "
|
||||||
|
f"{round(MIN_MEAN_CONFIDENCE * 100)}%"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
status = "good"
|
||||||
|
|
||||||
|
return {
|
||||||
|
"text": text,
|
||||||
|
"status": status,
|
||||||
|
"usable_for_summary": status == "good",
|
||||||
|
"line_count": len(lines),
|
||||||
|
"low_confidence_line_count": low_confidence_lines,
|
||||||
|
"mean_confidence": round(mean_confidence, 4),
|
||||||
|
"meaningful_chars": meaningful_chars,
|
||||||
|
"reasons": reasons,
|
||||||
|
"ocr_seconds": round(elapsed, 3),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _package_version(name: str) -> str | None:
|
||||||
|
try:
|
||||||
|
return importlib.metadata.version(name)
|
||||||
|
except importlib.metadata.PackageNotFoundError:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _extract(args) -> dict[str, Any]:
|
||||||
|
path = input_pdf(args.input)
|
||||||
|
page_count, page_sizes = _pdf_pages(path)
|
||||||
|
if page_count < 1:
|
||||||
|
raise ValueError("PDF 没有可执行 OCR 的页面")
|
||||||
|
requested_pages = parse_page_spec(args.pages, page_count)
|
||||||
|
if len(requested_pages) > MAX_PAGES_PER_CALL:
|
||||||
|
raise ValueError(
|
||||||
|
f"单次最多 OCR {MAX_PAGES_PER_CALL} 页,"
|
||||||
|
"请拆分 pages 后重试"
|
||||||
|
)
|
||||||
|
if args.start_offset > 0 and len(requested_pages) != 1:
|
||||||
|
raise ValueError("使用 start-offset 时 pages 必须只包含一页")
|
||||||
|
|
||||||
|
executable = _find_pdftoppm()
|
||||||
|
engine = _create_ocr_engine()
|
||||||
|
page_outputs: list[dict[str, Any]] = []
|
||||||
|
returned_chars = 0
|
||||||
|
next_page: int | None = None
|
||||||
|
next_offset = 0
|
||||||
|
remaining_pages: list[int] = []
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory(prefix="pdf-ocr-") as temp_dir_name:
|
||||||
|
temp_dir = Path(temp_dir_name)
|
||||||
|
for index, page_number in enumerate(requested_pages):
|
||||||
|
budget = args.max_chars - returned_chars
|
||||||
|
if budget <= 0:
|
||||||
|
next_page = page_number
|
||||||
|
remaining_pages = requested_pages[index:]
|
||||||
|
break
|
||||||
|
|
||||||
|
image_path, render_seconds = _render_page(
|
||||||
|
executable,
|
||||||
|
path,
|
||||||
|
page_number,
|
||||||
|
page_sizes[page_number],
|
||||||
|
args.dpi,
|
||||||
|
args.timeout,
|
||||||
|
temp_dir,
|
||||||
|
)
|
||||||
|
result = _ocr_page(engine, image_path)
|
||||||
|
full_text = result.pop("text")
|
||||||
|
offset = args.start_offset if index == 0 else 0
|
||||||
|
if offset > len(full_text):
|
||||||
|
raise ValueError(
|
||||||
|
f"start-offset 超过第 {page_number} 页 OCR "
|
||||||
|
f"文本长度 {len(full_text)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
usable = bool(result["usable_for_summary"])
|
||||||
|
if not usable:
|
||||||
|
page_text = ""
|
||||||
|
complete = True
|
||||||
|
else:
|
||||||
|
remaining_text = full_text[offset:]
|
||||||
|
page_text = remaining_text[:budget]
|
||||||
|
complete = len(page_text) == len(remaining_text)
|
||||||
|
|
||||||
|
page_outputs.append(
|
||||||
|
{
|
||||||
|
"page": page_number,
|
||||||
|
"text": page_text,
|
||||||
|
"char_count": len(full_text),
|
||||||
|
"offset_start": offset if usable else 0,
|
||||||
|
"offset_end": offset + len(page_text) if usable else 0,
|
||||||
|
"complete": complete,
|
||||||
|
"render_seconds": round(render_seconds, 3),
|
||||||
|
**result,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
returned_chars += len(page_text)
|
||||||
|
|
||||||
|
if not complete:
|
||||||
|
next_page = page_number
|
||||||
|
next_offset = offset + len(page_text)
|
||||||
|
remaining_pages = requested_pages[index + 1 :]
|
||||||
|
break
|
||||||
|
|
||||||
|
has_more = next_page is not None
|
||||||
|
all_processed = len(page_outputs) == len(requested_pages)
|
||||||
|
all_complete = all(page["complete"] for page in page_outputs)
|
||||||
|
all_usable = all(
|
||||||
|
page["usable_for_summary"] for page in page_outputs
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"path": str(path),
|
||||||
|
"page_count": page_count,
|
||||||
|
"engine": "rapidocr",
|
||||||
|
"engine_version": _package_version("rapidocr"),
|
||||||
|
"runtime": "onnxruntime",
|
||||||
|
"runtime_version": _package_version("onnxruntime"),
|
||||||
|
"offline": True,
|
||||||
|
"dpi": args.dpi,
|
||||||
|
"requested_pages": requested_pages,
|
||||||
|
"processed_pages": [page["page"] for page in page_outputs],
|
||||||
|
"returned_chars": returned_chars,
|
||||||
|
"pages": page_outputs,
|
||||||
|
"usable_for_summary": any(
|
||||||
|
page["usable_for_summary"] for page in page_outputs
|
||||||
|
),
|
||||||
|
"complete_ocr_coverage": (
|
||||||
|
all_processed and all_complete and all_usable
|
||||||
|
),
|
||||||
|
"needs_review": any(
|
||||||
|
not page["usable_for_summary"] for page in page_outputs
|
||||||
|
),
|
||||||
|
"has_more": has_more,
|
||||||
|
"next_page": next_page,
|
||||||
|
"next_offset": next_offset,
|
||||||
|
"remaining_pages": remaining_pages,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def main(argv: list[str] | None = None) -> int:
|
||||||
|
arguments = sys.argv[1:] if argv is None else argv
|
||||||
|
return run_cli(lambda: _extract(_parse_args(arguments)))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
raise SystemExit(main())
|
||||||
Loading…
Reference in New Issue
Block a user