wechat-robot-skills/skills/pptx/scripts/ocr_presentation.py

729 lines
22 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
from __future__ import annotations
import contextlib
import difflib
import importlib.metadata
import io
import logging
import os
import re
import shutil
import tempfile
import time
import unicodedata
from pathlib import Path
from typing import Any, Iterable
from _pptx_common import (
OOXML_PRESENTATION_SUFFIXES,
SkillArgumentParser,
find_program,
input_file,
run_cli,
run_program,
run_soffice_convert,
)
DEFAULT_DPI = 260
DEFAULT_MAX_CHARS = 24000
DEFAULT_TIMEOUT_SECONDS = 180
MAX_SLIDES_PER_CALL = 4
MAX_PIXELS_PER_SLIDE = 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 build_parser():
parser = SkillArgumentParser(
description="渲染指定演示文稿页面并用本地 OCR 提取图片文字。"
)
parser.add_argument("--input", required=True)
parser.add_argument(
"--slides",
required=True,
help="要识别的页码,例如 2 或 2,5-6;单次最多 4 页",
)
parser.add_argument(
"--start-offset",
type=int,
default=0,
help="续读单页图片文字时的字符偏移量",
)
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_TIMEOUT_SECONDS,
help=f"转换和单页渲染超时秒数,默认 {DEFAULT_TIMEOUT_SECONDS}",
)
return parser
def _parse_slide_spec(value: str, slide_count: int) -> list[int]:
if not value.strip():
raise ValueError("slides 不能为空")
slides: set[int] = set()
for raw_part in value.split(","):
part = raw_part.strip()
if not part:
continue
if "-" in part:
pieces = part.split("-", 1)
try:
start = int(pieces[0])
end = int(pieces[1])
except ValueError as exc:
raise ValueError(f"页码范围格式错误:{part}") from exc
if start > end:
raise ValueError(f"页码范围起始值不能大于结束值:{part}")
else:
try:
start = end = int(part)
except ValueError as exc:
raise ValueError(f"页码格式错误:{part}") from exc
if start < 1 or end > slide_count:
raise ValueError(f"页码必须在 1 到 {slide_count} 之间:{part}")
slides.update(range(start, end + 1))
if not slides:
raise ValueError("slides 不能为空")
return sorted(slides)
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 _comparison_key(value: str) -> str:
normalized = unicodedata.normalize("NFKC", value).casefold()
return "".join(character for character in normalized if character.isalnum())
def _iter_shapes(shapes: Any) -> Iterable[Any]:
from pptx.enum.shapes import MSO_SHAPE_TYPE
for shape in shapes:
yield shape
if shape.shape_type == MSO_SHAPE_TYPE.GROUP:
yield from _iter_shapes(shape.shapes)
def _shape_text_fragments(shape: Any) -> list[str]:
fragments: list[str] = []
if getattr(shape, "has_text_frame", False):
fragments.extend(
line
for line in _clean_text(shape.text_frame.text).splitlines()
if line
)
if getattr(shape, "has_table", False):
for row in shape.table.rows:
for cell in row.cells:
fragments.extend(
line
for line in _clean_text(cell.text).splitlines()
if line
)
return fragments
def _normalized_box(
shape: Any,
slide_width: int,
slide_height: int,
) -> tuple[float, float, float, float] | None:
try:
left = float(shape.left)
top = float(shape.top)
right = left + float(shape.width)
bottom = top + float(shape.height)
except (AttributeError, TypeError, ValueError):
return None
if slide_width <= 0 or slide_height <= 0:
return None
x0 = max(0.0, min(1.0, left / slide_width))
y0 = max(0.0, min(1.0, top / slide_height))
x1 = max(0.0, min(1.0, right / slide_width))
y1 = max(0.0, min(1.0, bottom / slide_height))
if x1 <= x0 or y1 <= y0:
return None
return (x0, y0, x1, y1)
def _contains_picture(shape: Any) -> bool:
from pptx.enum.shapes import MSO_SHAPE_TYPE
if shape.shape_type == MSO_SHAPE_TYPE.PICTURE:
return True
if shape.shape_type == MSO_SHAPE_TYPE.GROUP:
return any(_contains_picture(child) for child in shape.shapes)
return False
def _slide_profile(
slide: Any,
slide_width: int,
slide_height: int,
) -> dict[str, Any]:
from pptx.enum.shapes import MSO_SHAPE_TYPE
all_shapes = list(_iter_shapes(slide.shapes))
native_fragments: list[str] = []
picture_count = 0
chart_count = 0
for shape in all_shapes:
native_fragments.extend(_shape_text_fragments(shape))
if shape.shape_type == MSO_SHAPE_TYPE.PICTURE:
picture_count += 1
if getattr(shape, "has_chart", False):
chart_count += 1
unique_fragments = list(dict.fromkeys(native_fragments))
native_keys = [
key
for fragment in unique_fragments
if (key := _comparison_key(fragment))
]
native_boxes: list[dict[str, Any]] = []
picture_area = 0.0
for shape in slide.shapes:
box = _normalized_box(shape, slide_width, slide_height)
shape_fragments = _shape_text_fragments(shape)
if box and shape_fragments:
native_boxes.append(
{
"box": box,
"keys": [
key
for fragment in shape_fragments
if (key := _comparison_key(fragment))
],
}
)
if box and _contains_picture(shape):
picture_area += (box[2] - box[0]) * (box[3] - box[1])
return {
"picture_count": picture_count,
"chart_count": chart_count,
"image_area_ratio": round(min(1.0, picture_area), 4),
"native_text_char_count": sum(
1
for fragment in unique_fragments
for character in fragment
if character.isalnum()
),
"_native_keys": native_keys,
"_native_boxes": native_boxes,
}
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 _similar_to_any(
candidate: str,
references: list[str],
*,
threshold: float,
) -> bool:
if not candidate:
return False
for reference in references:
if not reference:
continue
if candidate == reference:
return True
shorter = min(len(candidate), len(reference))
longer = max(len(candidate), len(reference))
if shorter >= 3 and candidate in reference:
return True
if (
shorter >= 3
and reference in candidate
and longer <= round(shorter * 1.25)
):
return True
if shorter >= 3 and difflib.SequenceMatcher(
None,
candidate,
reference,
).ratio() >= threshold:
return True
return False
def _line_center(
box: list[list[float]] | None,
image_width: int,
image_height: int,
) -> tuple[float, float] | None:
if not box or image_width <= 0 or image_height <= 0:
return None
return (
sum(point[0] for point in box) / len(box) / image_width,
sum(point[1] for point in box) / len(box) / image_height,
)
def _line_is_native(
line: dict[str, Any],
profile: dict[str, Any],
image_width: int,
image_height: int,
) -> bool:
candidate = _comparison_key(line["text"])
if _similar_to_any(
candidate,
profile["_native_keys"],
threshold=0.82,
):
return True
center = _line_center(line["box"], image_width, image_height)
if center is None:
return False
x, y = center
padding = 0.01
for native_box in profile["_native_boxes"]:
x0, y0, x1, y1 = native_box["box"]
if (
x0 - padding <= x <= x1 + padding
and y0 - padding <= y <= y1 + padding
and _similar_to_any(
candidate,
native_box["keys"],
threshold=0.68,
)
):
return True
return False
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 _ocr_slide(
engine: Any,
image_path: Path,
profile: dict[str, Any],
) -> dict[str, Any]:
from PIL import Image
with Image.open(image_path) as image:
image_width, image_height = image.size
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
raw_lines = _ordered_lines(result)
image_lines: list[dict[str, Any]] = []
seen: set[str] = set()
filtered_native = 0
filtered_duplicates = 0
for line in raw_lines:
if _line_is_native(line, profile, image_width, image_height):
filtered_native += 1
continue
key = _comparison_key(line["text"])
if key and key in seen:
filtered_duplicates += 1
continue
if key:
seen.add(key)
image_lines.append(line)
text = "\n".join(line["text"] for line in image_lines)
weighted_chars = [
max(1, sum(1 for character in line["text"] if not character.isspace()))
for line in image_lines
]
total_weight = sum(weighted_chars)
mean_confidence = (
sum(
line["confidence"] * weight
for line, weight in zip(image_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 image_lines
if line["confidence"] < MIN_MEAN_CONFIDENCE
)
reasons: list[str] = []
if not text:
status = "no_image_text"
reasons.append("未识别到原生文本之外的图片文字")
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",
"needs_review": status in {"sparse", "low_confidence"},
"raw_ocr_line_count": len(raw_lines),
"image_line_count": len(image_lines),
"filtered_native_line_count": filtered_native,
"filtered_duplicate_line_count": filtered_duplicates,
"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 _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("LibreOffice 生成了加密 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 _render_slide(
pdf_path: Path,
slide_number: int,
page_size: tuple[float, float],
dpi: int,
timeout: int,
temp_dir: Path,
) -> tuple[Path, float]:
width_points, height_points = page_size
estimated_pixels = (
width_points * dpi / 72.0
* height_points * dpi / 72.0
)
if estimated_pixels > MAX_PIXELS_PER_SLIDE:
raise ValueError(
f"第 {slide_number} 页按 {dpi} DPI 渲染预计超过 "
f"{MAX_PIXELS_PER_SLIDE} 像素,请降低 dpi"
)
prefix = temp_dir / f"slide-{slide_number:04d}"
output = prefix.with_suffix(".png")
started = time.monotonic()
run_program(
[
find_program("pdftoppm"),
"-f",
str(slide_number),
"-l",
str(slide_number),
"-singlefile",
"-png",
"-r",
str(dpi),
str(pdf_path),
str(prefix),
],
timeout=timeout,
)
elapsed = time.monotonic() - started
if not output.is_file() or output.stat().st_size <= 0:
raise RuntimeError(f"第 {slide_number} 页没有生成有效 PNG")
return output, elapsed
def _package_version(name: str) -> str | None:
try:
return importlib.metadata.version(name)
except importlib.metadata.PackageNotFoundError:
return None
def main() -> dict[str, Any]:
from pptx import Presentation
args = build_parser().parse_args()
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 秒之间")
source = input_file(args.input, OOXML_PRESENTATION_SUFFIXES)
presentation = Presentation(str(source))
slide_count = len(presentation.slides)
if slide_count < 1:
raise ValueError("演示文稿没有可执行 OCR 的页面")
requested_slides = _parse_slide_spec(args.slides, slide_count)
if len(requested_slides) > MAX_SLIDES_PER_CALL:
raise ValueError(
f"单次最多 OCR {MAX_SLIDES_PER_CALL} 页,请拆分 slides 后重试"
)
if args.start_offset > 0 and len(requested_slides) != 1:
raise ValueError("使用 start-offset 时 slides 必须只包含一页")
slide_width, slide_height = presentation.slide_width, presentation.slide_height
if slide_width is None or slide_height is None:
raise ValueError("演示文稿缺少页面尺寸")
profiles = {
slide_number: _slide_profile(
presentation.slides[slide_number - 1],
slide_width,
slide_height,
)
for slide_number in requested_slides
}
engine = _create_ocr_engine()
page_outputs: list[dict[str, Any]] = []
returned_chars = 0
next_slide: int | None = None
next_offset = 0
remaining_slides: list[int] = []
office_output = {"stdout": "", "stderr": ""}
with tempfile.TemporaryDirectory(prefix="pptx-ocr-") as temp_name:
temp_dir = Path(temp_name)
staged_input = temp_dir / f"presentation{source.suffix.lower()}"
shutil.copy2(source, staged_input)
pdf_path, office_output = run_soffice_convert(
staged_input,
target_format="pdf",
output_dir=temp_dir / "pdf",
timeout=args.timeout,
)
pdf_page_count, page_sizes = _pdf_pages(pdf_path)
if pdf_page_count != slide_count:
raise RuntimeError(
f"演示文稿有 {slide_count} 页,但渲染结果有 "
f"{pdf_page_count} 页"
)
for index, slide_number in enumerate(requested_slides):
budget = args.max_chars - returned_chars
if budget <= 0:
next_slide = slide_number
remaining_slides = requested_slides[index:]
break
image_path, render_seconds = _render_slide(
pdf_path,
slide_number,
page_sizes[slide_number],
args.dpi,
args.timeout,
temp_dir,
)
result = _ocr_slide(
engine,
image_path,
profiles[slide_number],
)
full_text = result.pop("text")
offset = args.start_offset if index == 0 else 0
if offset > len(full_text):
raise ValueError(
f"start-offset 超过第 {slide_number} 页图片文字长度 "
f"{len(full_text)}"
)
usable = bool(result["usable_for_summary"])
if not usable:
slide_text = ""
complete = True
else:
remaining_text = full_text[offset:]
slide_text = remaining_text[:budget]
complete = len(slide_text) == len(remaining_text)
profile = profiles[slide_number]
page_outputs.append(
{
"slide": slide_number,
"text": slide_text,
"char_count": len(full_text),
"offset_start": offset if usable else 0,
"offset_end": offset + len(slide_text) if usable else 0,
"complete": complete,
"render_seconds": round(render_seconds, 3),
"picture_count": profile["picture_count"],
"chart_count": profile["chart_count"],
"image_area_ratio": profile["image_area_ratio"],
"native_text_char_count": profile[
"native_text_char_count"
],
**result,
}
)
returned_chars += len(slide_text)
if not complete:
next_slide = slide_number
next_offset = offset + len(slide_text)
remaining_slides = requested_slides[index + 1 :]
break
all_processed = len(page_outputs) == len(requested_slides)
all_complete = all(page["complete"] for page in page_outputs)
all_safe = all(
page["status"] in {"good", "no_image_text"}
for page in page_outputs
)
return {
"source": str(source),
"slide_count": slide_count,
"engine": "rapidocr",
"engine_version": _package_version("rapidocr"),
"runtime": "onnxruntime",
"runtime_version": _package_version("onnxruntime"),
"offline": True,
"dpi": args.dpi,
"requested_slides": requested_slides,
"processed_slides": [page["slide"] for page in page_outputs],
"returned_chars": returned_chars,
"slides": 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_safe
),
"needs_review": any(page["needs_review"] for page in page_outputs),
"has_more": next_slide is not None,
"next_slide": next_slide,
"next_offset": next_offset,
"remaining_slides": remaining_slides,
"office_stdout": office_output["stdout"],
"office_stderr": office_output["stderr"],
}
if __name__ == "__main__":
raise SystemExit(run_cli(main))