Files
my_wiki/scripts/paddleocr_v2_ocr.py

228 lines
9.2 KiB
Python

#!/usr/bin/env python3
"""
PaddleOCR v2 jobs async API OCR tool (user-verified, PP-OCRv6).
Usage:
python paddleocr_v2_ocr.py <file_path_or_url> [--model PP-OCRv6] [--out DIR]
Token: read from env PADDLEOCR_ACCESS_TOKEN (already in hermes-home/.env),
fallback to hardcoded token below.
Reference: official PaddleOCR v2 jobs API example (2026-08).
Verified working 2026-08-07 with Chinese financial text image.
"""
import json
import os
import sys
import time
import requests
JOB_URL = "https://paddleocr.aistudio-app.com/api/v2/ocr/jobs"
def _load_token() -> str:
"""Load PADDLEOCR_ACCESS_TOKEN from env, then hermes-home/.env as fallback."""
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()
MODEL = "PP-OCRv6"
optional_payload = {
"useDocOrientationClassify": False,
"useDocUnwarping": False,
"useTextlineOrientation": False,
}
def _html_table_to_markdown(html: str) -> str:
"""Convert a simple <table> HTML string to a markdown table."""
if not html or "<table" not in html:
return html
import re
rows = re.findall(r"<tr>(.*?)</tr>", html, re.S)
md_lines = []
for i, row in enumerate(rows):
cells = re.findall(r"<t[dh]>(.*?)</t[dh]>", row, re.S)
cells = [c.strip() for c in cells]
md_lines.append("| " + " | ".join(cells) + " |")
if i == 0:
md_lines.append("| " + " | ".join(["---"] * len(cells)) + " |")
return "\n".join(md_lines)
def run(file_path: str, model: str = MODEL, out_dir: str = "output"):
headers = {"Authorization": f"bearer {TOKEN}"}
print(f"Processing file: {file_path}")
if file_path.startswith("http"):
headers["Content-Type"] = "application/json"
payload = {"fileUrl": file_path, "model": model, "optionalPayload": optional_payload}
job_response = requests.post(JOB_URL, json=payload, headers=headers)
else:
if not os.path.exists(file_path):
print(f"Error: File not found at {file_path}")
sys.exit(1)
data = {"model": model, "optionalPayload": json.dumps(optional_payload)}
with open(file_path, "rb") as f:
files = {"file": f}
job_response = requests.post(JOB_URL, headers=headers, data=data, files=files)
print(f"Response status: {job_response.status_code}")
if job_response.status_code != 200:
print(f"Response content: {job_response.text}")
assert job_response.status_code == 200
jobId = job_response.json()["data"]["jobId"]
print(f"Job submitted successfully. job id: {jobId}")
print("Start polling for results")
jsonl_url = ""
while True:
r = requests.get(f"{JOB_URL}/{jobId}", headers=headers)
assert r.status_code == 200
state = r.json()["data"]["state"]
if state == 'pending':
print("The current status of the job is pending")
elif state == 'running':
try:
tp = r.json()['data']['extractProgress']['totalPages']
ep = r.json()['data']['extractProgress']['extractedPages']
print(f"running, total pages: {tp}, extracted pages: {ep}")
except KeyError:
print("running...")
elif state == 'done':
ep = r.json()['data']['extractProgress']['extractedPages']
st = r.json()['data']['extractProgress']['startTime']
et = r.json()['data']['extractProgress']['endTime']
print(f"Job completed, pages: {ep}, {st} -> {et}")
jsonl_url = r.json()['data']['resultUrl']['jsonUrl']
break
elif state == "failed":
print(f"Job failed: {r.json()['data']['errorMsg']}")
sys.exit(1)
time.sleep(5)
if jsonl_url:
jr = requests.get(jsonl_url)
jr.raise_for_status()
lines = jr.text.strip().split('\n')
os.makedirs(out_dir, exist_ok=True)
page_num = 0
for line in lines:
line = line.strip()
if not line:
continue
obj = json.loads(line)
result = obj.get("result") or {}
page_num += 1
print(f"--- page {page_num} ---")
# Style A: PP-OCRv5/v6 (ocrResults)
if isinstance(result, dict) and "ocrResults" in result:
for res in result.get("ocrResults", []):
pruned = res.get("prunedResult", {})
texts = pruned.get("rec_texts", [])
scores = pruned.get("rec_scores", [])
for t, s in zip(texts, scores):
print(f" [{s:.2f}] {t}")
img_url = res.get("ocrImage")
if img_url:
ir = requests.get(img_url)
if ir.status_code == 200:
fn = os.path.join(out_dir, f"img_output_{page_num}.jpg")
with open(fn, "wb") as f:
f.write(ir.content)
print(f"Image saved to: {fn}")
# Style B: PaddleOCR-VL models (layoutParsingResults)
elif isinstance(result, dict) and "layoutParsingResults" in result:
for res in result.get("layoutParsingResults", []):
md = res.get("markdown") or {}
md_text = md.get("text")
if md_text:
md_path = os.path.join(out_dir, f"doc_{page_num}.md")
with open(md_path, "w", encoding="utf-8") as f:
f.write(md_text)
print(f"Markdown saved to: {md_path}")
print("--- markdown content ---")
print(md_text)
for img_rel, img_url in (md.get("images") or {}).items():
try:
ir = requests.get(img_url)
if ir.status_code == 200:
full = os.path.join(
out_dir, img_rel.replace("\\", "/")
)
os.makedirs(
os.path.dirname(full), exist_ok=True
)
with open(full, "wb") as f:
f.write(ir.content)
print(f"Markdown image saved to: {full}")
except Exception as e:
print(f"Markdown image download failed: {e}")
else:
# Fallback: block-level parsing list
pruned = res.get("prunedResult", {})
for blk in (pruned.get("parsing_res_list") or []):
label = blk.get("block_label", "?")
content = blk.get("block_content", "")
if label == "table":
print(" [table] markdown:")
print(_html_table_to_markdown(content))
else:
print(f" [{label}] {content}")
for img_name, img_url in (res.get("outputImages") or {}).items():
try:
ir = requests.get(img_url)
if ir.status_code == 200:
fn = os.path.join(
out_dir, f"{img_name}_{page_num}.jpg"
)
with open(fn, "wb") as f:
f.write(ir.content)
print(f"Image saved to: {fn}")
except Exception as e:
print(f"Image download failed: {e}")
else:
print(" (unknown result schema; raw json saved below)")
# Save raw result per page
raw_path = os.path.join(out_dir, f"page_{page_num}.json")
with open(raw_path, "w", encoding="utf-8") as f:
f.write(json.dumps(obj, ensure_ascii=False, indent=2))
print(f"Raw saved to: {raw_path}")
if __name__ == "__main__":
if len(sys.argv) < 2:
print(__doc__)
sys.exit(1)
file_arg = sys.argv[1]
model_arg = MODEL
out_arg = "output"
i = 2
while i < len(sys.argv):
if sys.argv[i] == "--model" and i + 1 < len(sys.argv):
model_arg = sys.argv[i + 1]
i += 2
elif sys.argv[i] == "--out" and i + 1 < len(sys.argv):
out_arg = sys.argv[i + 1]
i += 2
else:
i += 1
run(file_arg, model_arg, out_arg)