vent_agent_2.0/algorithms/main_fan/document_worker.py

"""隔离 PDFium 原生解析；输出有限页数及分辨率的文本和 PNG。"""
import argparse
import json
from pathlib import Path
import re


def text_is_readable(text):
    """拒绝空层、零散水印及大量损坏字符；完整文字页可绕过图像模型。"""
    compact = re.sub(r"\s+", "", text)
    if len(compact) < 100:
        return False
    meaningful = sum(c.isalpha() or c.isdigit() for c in compact)
    broken = sum(c in "\ufffd\u25a1\u25a0" or ord(c) < 32 for c in compact)
    return meaningful / len(compact) >= .5 and broken / len(compact) < .03


def main():
    import pypdfium2 as pdfium
    parser = argparse.ArgumentParser()
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--max-pages", type=int, default=32)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    with pdfium.PdfDocument(args.source) as doc:
        if len(doc) > args.max_pages:
            raise ValueError("PDF 页数超过识别上限，请拆分报告")
        rows = []
        for index in range(len(doc)):
            page = doc[index]
            try:
                textpage = page.get_textpage()
                try:
                    content = textpage.get_text_range()
                finally:
                    textpage.close()
                coverage = 0.0
                for obj in page.get_objects():
                    if obj.type == pdfium.raw.FPDF_PAGEOBJ_IMAGE:
                        x0, y0, x1, y1 = obj.get_bounds()
                        width, height = page.get_size()
                        coverage += abs((x1-x0)*(y1-y0))/(width*height)
                readable = text_is_readable(content)
                # 大图仅带少量文字、以及曲线/图页仍走视觉，避免忽略正文图像。
                visual = not readable or coverage >= .5 and len(content.strip()) < 400 or bool(re.search("性能曲线|特性曲线", content))
                image = args.output / f"page-{index+1}.png"
                if visual:
                    bitmap = page.render(scale=min(3, 2200/max(page.get_size())))
                    try:
                        bitmap.to_pil().save(image)
                    finally:
                        bitmap.close()
                rows.append({"page": index+1, "text": content,
                             "image": image.name if visual else None,
                             "route": "vision" if visual else "text",
                             "text_chars": len(content), "image_coverage": round(min(coverage,1),4)})
            finally:
                page.close()
    (args.output / "pages.json").write_text(json.dumps(rows, ensure_ascii=False), encoding="utf-8")


if __name__ == "__main__":
    main()