182 lines
6.9 KiB
Python
182 lines
6.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
书籍扫描件批量 OCR -> 单本 Markdown 全文版（PaddleOCR-VL-1.6 整页 Markdown）。
|
||
|
||
背景：PaddleOCR v2 jobs API 单文件 100 页截断，整本书需按 100 页拆片提交；
|
||
VL-1.6 模型返回每页 layoutParsingResults[].markdown.text（整页 Markdown，
|
||
含表格/版面还原）。本脚本把全书合并为一份 .md 存档（raw/书籍/_ocr-md/）。
|
||
|
||
Usage:
|
||
python ocr_book_to_md.py <input.pdf> <output.md> [--pages 100] [--keep-chunks]
|
||
|
||
流程：
|
||
1. fitz 拆片（默认每 100 页一片，临时 PDF 存 <out_dir>/_chunks/）
|
||
2. 每片提交 PaddleOCR-VL-1.6 job 并轮询（失败自动重试 2 次）
|
||
3. done 后下载 jsonl 缓存到 _chunks/chunk_XX.jsonl（已存在则跳过=断点续跑）
|
||
4. 解析 jsonl 中 markdown 文本，按页序合并为 output.md
|
||
"""
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
import requests
|
||
import fitz # PyMuPDF
|
||
|
||
JOB_URL = "https://paddleocr.aistudio-app.com/api/v2/ocr/jobs"
|
||
MODEL = "PaddleOCR-VL-1.6"
|
||
OPTIONAL = {"useDocOrientationClassify": False, "useDocUnwarping": False, "useTextlineOrientation": False}
|
||
|
||
|
||
def load_token():
|
||
tok = os.environ.get("PADDLEOCR_ACCESS_TOKEN", "").strip()
|
||
if tok:
|
||
return tok
|
||
env_path = os.path.join(
|
||
os.path.expanduser("~"), "AppData", "Local",
|
||
"Hermes Agent CN Desktop", "data", "hermes-home", ".env",
|
||
)
|
||
try:
|
||
with open(env_path, encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line.startswith("PADDLEOCR_ACCESS_TOKEN="):
|
||
return line.split("=", 1)[1].strip()
|
||
except OSError:
|
||
pass
|
||
return ""
|
||
|
||
|
||
TOKEN = load_token()
|
||
if not TOKEN:
|
||
print("FATAL: PADDLEOCR_ACCESS_TOKEN not found")
|
||
sys.exit(1)
|
||
|
||
|
||
def submit_and_wait(pdf_path: str, headers: dict, retries: int = 2):
|
||
"""提交一片 PDF，轮询至 done，返回 jsonl 文本。失败重试。"""
|
||
data = {"model": MODEL, "optionalPayload": json.dumps(OPTIONAL)}
|
||
for attempt in range(1, retries + 2):
|
||
try:
|
||
with open(pdf_path, "rb") as f:
|
||
resp = requests.post(JOB_URL, headers=headers, data=data, files={"file": f}, timeout=120)
|
||
if resp.status_code != 200:
|
||
print(f" submit HTTP {resp.status_code}: {resp.text[:300]}")
|
||
raise RuntimeError(f"submit failed: {resp.status_code}")
|
||
job_id = resp.json()["data"]["jobId"]
|
||
print(f" job={job_id} attempt={attempt}")
|
||
while True:
|
||
time.sleep(8)
|
||
r = requests.get(f"{JOB_URL}/{job_id}", headers=headers, timeout=60)
|
||
st = r.json()["data"]["state"]
|
||
if st == "running":
|
||
try:
|
||
ep = r.json()["data"]["extractProgress"]["extractedPages"]
|
||
tp = r.json()["data"]["extractProgress"]["totalPages"]
|
||
print(f" running {ep}/{tp}")
|
||
except KeyError:
|
||
print(" running...")
|
||
elif st == "done":
|
||
jurl = r.json()["data"]["resultUrl"]["jsonUrl"]
|
||
jr = requests.get(jurl, timeout=120)
|
||
jr.raise_for_status()
|
||
print(f" done ({len(jr.text)} bytes jsonl)")
|
||
return jr.text
|
||
elif st == "failed":
|
||
msg = r.json()["data"].get("errorMsg", "?")
|
||
print(f" FAILED: {msg}")
|
||
raise RuntimeError(f"job failed: {msg}")
|
||
except (RuntimeError, requests.RequestException) as e:
|
||
print(f" attempt {attempt} error: {e}")
|
||
if attempt <= retries:
|
||
print(" retrying in 20s...")
|
||
time.sleep(20)
|
||
else:
|
||
raise
|
||
|
||
|
||
def parse_jsonl(jsonl_text: str):
|
||
"""解析 jsonl -> list[page_md]（按行序）。"""
|
||
pages = []
|
||
for line in jsonl_text.strip().split("\n"):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
obj = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
result = obj.get("result") or {}
|
||
# VL 模型 jsonl 分块返回：每行含 N 个 layoutParsingResults（=N 页），
|
||
# 每个 lpr 的 markdown.text 即一页；dataInfo.numPages 佐证块内页数。
|
||
lpr = result.get("layoutParsingResults") or []
|
||
for res in lpr:
|
||
md = res.get("markdown") or {}
|
||
t = md.get("text")
|
||
if t:
|
||
pages.append(t)
|
||
return pages
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) < 3:
|
||
print(__doc__)
|
||
sys.exit(1)
|
||
pdf_path, out_md = sys.argv[1], sys.argv[2]
|
||
pages_per_chunk = 100
|
||
keep_chunks = False
|
||
if "--pages" in sys.argv:
|
||
pages_per_chunk = int(sys.argv[sys.argv.index("--pages") + 1])
|
||
if "--keep-chunks" in sys.argv:
|
||
keep_chunks = True
|
||
|
||
base = os.path.dirname(os.path.abspath(out_md))
|
||
chunks_dir = os.path.join(base, "_chunks")
|
||
os.makedirs(chunks_dir, exist_ok=True)
|
||
headers = {"Authorization": f"bearer {TOKEN}"}
|
||
|
||
doc = fitz.open(pdf_path)
|
||
total = doc.page_count
|
||
print(f"PDF pages: {total}, chunks of {pages_per_chunk}")
|
||
|
||
all_pages = []
|
||
chunk_idx = 0
|
||
for start in range(0, total, pages_per_chunk):
|
||
end = min(start + pages_per_chunk, total)
|
||
chunk_idx += 1
|
||
chunk_pdf = os.path.join(chunks_dir, f"chunk_{chunk_idx:02d}_{start+1}-{end}.pdf")
|
||
jsonl_cache = chunk_pdf.replace(".pdf", ".jsonl")
|
||
print(f"\n=== chunk {chunk_idx}: pages {start+1}-{end} ===")
|
||
if os.path.exists(jsonl_cache):
|
||
print(f" cached jsonl exists, skip submit: {os.path.basename(jsonl_cache)}")
|
||
with open(jsonl_cache, encoding="utf-8") as f:
|
||
jsonl_text = f.read()
|
||
else:
|
||
if not os.path.exists(chunk_pdf):
|
||
sub = fitz.open()
|
||
sub.insert_pdf(doc, from_page=start, to_page=end - 1)
|
||
sub.save(chunk_pdf, garbage=3)
|
||
sub.close()
|
||
print(f" split -> {os.path.basename(chunk_pdf)}")
|
||
jsonl_text = submit_and_wait(chunk_pdf, headers)
|
||
with open(jsonl_cache, "w", encoding="utf-8") as f:
|
||
f.write(jsonl_text)
|
||
pages = parse_jsonl(jsonl_text)
|
||
print(f" parsed {len(pages)} pages")
|
||
all_pages.extend(pages)
|
||
if not keep_chunks:
|
||
try:
|
||
os.remove(chunk_pdf)
|
||
except OSError:
|
||
pass
|
||
doc.close()
|
||
|
||
with open(out_md, "w", encoding="utf-8") as f:
|
||
f.write(f"<!-- OCR: {os.path.basename(pdf_path)} | {MODEL} | {total} pages | {time.strftime('%Y-%m-%d %H:%M')} -->\n\n")
|
||
for i, pg in enumerate(all_pages, 1):
|
||
f.write(f"\n\n<!-- page {i} -->\n\n{pg.strip()}")
|
||
print(f"\nDONE: {out_md} ({len(all_pages)} pages, {os.path.getsize(out_md)} bytes)")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|