79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
用于 mindmap-ai v2：从 PDF 中提取分页文本。
|
|
|
|
注意：
|
|
- 本脚本仅做“文本抽取”，不做大纲推断（大纲交给 TS 侧 + LLM/规则）。
|
|
- 输出为 JSON，便于 Next.js API 通过 child_process 调用并解析。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
|
|
def _load_reader(path: str):
|
|
try:
|
|
# pypdf (推荐)
|
|
from pypdf import PdfReader # type: ignore
|
|
return PdfReader(path)
|
|
except Exception:
|
|
# PyPDF2 兜底
|
|
from PyPDF2 import PdfReader # type: ignore
|
|
return PdfReader(path)
|
|
|
|
|
|
def _extract_pages(reader, max_pages: Optional[int] = None) -> List[Dict[str, Any]]:
|
|
pages: List[Dict[str, Any]] = []
|
|
total = len(getattr(reader, "pages", []) or [])
|
|
limit = total if max_pages is None else min(total, max_pages)
|
|
|
|
for i in range(limit):
|
|
try:
|
|
page = reader.pages[i]
|
|
text = page.extract_text() or ""
|
|
except Exception:
|
|
text = ""
|
|
# 清理一些不可见字符，避免 JSON/解析异常
|
|
text = text.replace("\x00", "").strip()
|
|
pages.append({"page": i + 1, "text": text})
|
|
return pages
|
|
|
|
|
|
def main(argv: List[str]) -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input", required=True, help="PDF 文件路径")
|
|
parser.add_argument("--max-pages", type=int, default=50, help="最多提取页数")
|
|
args = parser.parse_args(argv)
|
|
|
|
reader = _load_reader(args.input)
|
|
meta_title = None
|
|
try:
|
|
info = getattr(reader, "metadata", None)
|
|
meta_title = getattr(info, "title", None) if info else None
|
|
except Exception:
|
|
meta_title = None
|
|
|
|
pages = _extract_pages(reader, max_pages=args.max_pages)
|
|
payload: Dict[str, Any] = {
|
|
"meta": {"title": meta_title},
|
|
"pages": pages,
|
|
"totalPages": len(getattr(reader, "pages", []) or []),
|
|
}
|
|
# Windows 控制台默认编码可能是 gbk，直接写入会导致 UnicodeEncodeError
|
|
try:
|
|
sys.stdout.reconfigure(encoding="utf-8") # py3.7+
|
|
except Exception:
|
|
pass
|
|
sys.stdout.write(json.dumps(payload, ensure_ascii=False))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main(sys.argv[1:]))
|