457 lines
14 KiB
Python
457 lines
14 KiB
Python
#!/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())
|