299 lines
11 KiB
Python
299 lines
11 KiB
Python
"""
|
||
ocr_confidence.py
|
||
-----------------
|
||
Reads an existing olmOCR .md output file and scores each page for transcription
|
||
confidence using the DeepInfra API. Does NOT re-run OCR.
|
||
|
||
Requirements:
|
||
pip install requests
|
||
|
||
Usage:
|
||
python ocr_confidence.py --api_key "YOUR_KEY"
|
||
python ocr_confidence.py --api_key "YOUR_KEY" --input "path/to/file.md" --output "path/to/out.json"
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
import warnings
|
||
from pathlib import Path
|
||
|
||
import requests
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Load .env from project root (if present) — no external dependencies needed
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _load_dotenv() -> None:
|
||
env_path = Path(__file__).resolve().parent.parent / ".env"
|
||
if not env_path.exists():
|
||
return
|
||
import os
|
||
with open(env_path, encoding="utf-8") as fh:
|
||
for line in fh:
|
||
line = line.strip()
|
||
if not line or line.startswith("#") or "=" not in line:
|
||
continue
|
||
key, _, value = line.partition("=")
|
||
os.environ.setdefault(key.strip(), value.strip())
|
||
|
||
_load_dotenv()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration — mirrors ocr_wardiaries.py
|
||
# ---------------------------------------------------------------------------
|
||
|
||
DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
|
||
# olmOCR is a vision model — use a chat LLM for text-based confidence review
|
||
MODEL = "google/gemma-3-27b-it"
|
||
MAX_RETRIES = 3
|
||
RETRY_DELAY = 5 # seconds between retries
|
||
|
||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||
DEFAULT_INPUT = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_olmocr.md"
|
||
DEFAULT_OUTPUT = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_confidence.json"
|
||
|
||
CONFIDENCE_PROMPT = (
|
||
"You are reviewing an OCR transcription of a WWII war diary. Your job is to identify words that may have been misread by the OCR scanner — not to correct the soldier's original spelling or interpret abbreviations. Rate OCR accuracy only. Do not suggest what words "should" be. Respond in JSON only:
|
||
{"score": 7, "uncertain_words": ["Loon", "Fme"], "notes": "Possible OCR misread in line 3"}
|
||
)
|
||
|
||
ERROR_RESULT = {"score": 0, "uncertain_words": [], "notes": "API error"}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Parse olmOCR markdown into {page_num: text} dict
|
||
# Same logic as parseOCRByPage() in p44-ocr-viewer.html
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def parse_ocr_by_page(raw_text: str) -> dict[int, str]:
|
||
"""Split olmOCR markdown on '## Page N' headings into a page-keyed dict."""
|
||
pages: dict[int, str] = {}
|
||
lines = raw_text.split("\n")
|
||
page_num: int | None = None
|
||
buf: list[str] = []
|
||
|
||
for line in lines:
|
||
m_head = _PAGE_HEADING.match(line)
|
||
if m_head:
|
||
if page_num is not None:
|
||
pages[page_num] = "\n".join(buf)
|
||
page_num = int(m_head.group(1))
|
||
buf = [line] # keep the heading at the top
|
||
elif page_num is not None:
|
||
buf.append(line)
|
||
|
||
if page_num is not None:
|
||
pages[page_num] = "\n".join(buf)
|
||
|
||
return pages
|
||
|
||
|
||
import re as _re
|
||
_PAGE_HEADING = _re.compile(r"^## Page (\d+)\s*$")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# API call
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def score_page(page_text: str, page_num: int, api_key: str) -> dict:
|
||
"""
|
||
Send one page's OCR text to DeepInfra for confidence scoring.
|
||
Returns a dict with keys: score, uncertain_words, notes.
|
||
On failure returns ERROR_RESULT.
|
||
"""
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {api_key}",
|
||
}
|
||
payload = {
|
||
"model": MODEL,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
f"{CONFIDENCE_PROMPT}\n\n"
|
||
f"--- BEGIN OCR TEXT (page {page_num}) ---\n"
|
||
f"{page_text}\n"
|
||
f"--- END OCR TEXT ---"
|
||
),
|
||
}
|
||
],
|
||
"max_tokens": 256,
|
||
"temperature": 0.0,
|
||
}
|
||
|
||
for attempt in range(1, MAX_RETRIES + 1):
|
||
try:
|
||
response = requests.post(
|
||
DEEPINFRA_API_URL, headers=headers, json=payload, timeout=60
|
||
)
|
||
response.raise_for_status()
|
||
raw_content = response.json()["choices"][0]["message"]["content"].strip()
|
||
|
||
# Strip markdown code fences if the model wrapped the JSON
|
||
if raw_content.startswith("```"):
|
||
raw_content = _re.sub(r"^```[a-z]*\n?", "", raw_content)
|
||
raw_content = _re.sub(r"\n?```$", "", raw_content)
|
||
|
||
result = json.loads(raw_content)
|
||
|
||
# Validate expected keys; fill missing ones with defaults
|
||
score = int(result.get("score", 0))
|
||
uncertain_words = result.get("uncertain_words", [])
|
||
notes = result.get("notes", "")
|
||
|
||
if not isinstance(uncertain_words, list):
|
||
uncertain_words = []
|
||
|
||
return {"score": score, "uncertain_words": uncertain_words, "notes": notes}
|
||
|
||
except requests.exceptions.HTTPError as exc:
|
||
warnings.warn(f"Page {page_num}: HTTP error on attempt {attempt}/{MAX_RETRIES}: {exc}")
|
||
if attempt < MAX_RETRIES:
|
||
time.sleep(RETRY_DELAY)
|
||
|
||
except json.JSONDecodeError as exc:
|
||
warnings.warn(
|
||
f"Page {page_num}: Malformed JSON from API on attempt {attempt}/{MAX_RETRIES}: {exc}"
|
||
)
|
||
if attempt < MAX_RETRIES:
|
||
time.sleep(RETRY_DELAY)
|
||
|
||
except Exception as exc: # noqa: BLE001
|
||
warnings.warn(f"Page {page_num}: Unexpected error on attempt {attempt}/{MAX_RETRIES}: {exc}")
|
||
if attempt < MAX_RETRIES:
|
||
time.sleep(RETRY_DELAY)
|
||
|
||
warnings.warn(f"Page {page_num}: All {MAX_RETRIES} attempts failed — storing error result.")
|
||
return dict(ERROR_RESULT) # return a fresh copy
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Summary helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def print_summary(results: dict) -> None:
|
||
high = sum(1 for v in results.values() if v["score"] >= 8)
|
||
med = sum(1 for v in results.values() if 5 <= v["score"] <= 7)
|
||
low = sum(1 for v in results.values() if 1 <= v["score"] <= 4)
|
||
err = sum(1 for v in results.values() if v["score"] == 0)
|
||
|
||
total = len(results)
|
||
print("\n" + "=" * 50)
|
||
print(f"CONFIDENCE SCORING COMPLETE — {total} pages scored")
|
||
print("=" * 50)
|
||
print(f" High (8-10) : {high:>4} ({high/total*100:.1f}%)" if total else " High (8-10) : 0")
|
||
print(f" Medium (5-7) : {med:>4} ({med/total*100:.1f}%)" if total else " Medium (5-7) : 0")
|
||
print(f" Low (1-4) : {low:>4} ({low/total*100:.1f}%)" if total else " Low (1-4) : 0")
|
||
if err:
|
||
print(f" API errors : {err:>4}")
|
||
print("=" * 50)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(
|
||
description="Score OCR confidence for each page of an olmOCR .md output file."
|
||
)
|
||
parser.add_argument(
|
||
"--api_key", default=None,
|
||
help="DeepInfra API key (defaults to DEEPINFRA_API_KEY env var)."
|
||
)
|
||
parser.add_argument(
|
||
"--input", default=str(DEFAULT_INPUT),
|
||
help=f"Path to olmOCR .md file. Default: {DEFAULT_INPUT}"
|
||
)
|
||
parser.add_argument(
|
||
"--output", default=str(DEFAULT_OUTPUT),
|
||
help=f"Path to write confidence JSON. Default: {DEFAULT_OUTPUT}"
|
||
)
|
||
parser.add_argument(
|
||
"--delay", type=float, default=0.5,
|
||
help="Seconds to pause between API calls (default: 0.5)."
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
import os
|
||
api_key = args.api_key or os.environ.get("DEEPINFRA_API_KEY", "")
|
||
if not api_key:
|
||
print("ERROR: No API key provided. Use --api_key or set DEEPINFRA_API_KEY in .env", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
input_path = Path(args.input)
|
||
output_path = Path(args.output)
|
||
|
||
# ── Read source file ──────────────────────────────────────────────────────
|
||
if not input_path.exists():
|
||
print(f"ERROR: Input file not found: {input_path}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
print(f"Reading OCR source: {input_path}")
|
||
raw_text = input_path.read_text(encoding="utf-8")
|
||
|
||
pages = parse_ocr_by_page(raw_text)
|
||
if not pages:
|
||
print("ERROR: No '## Page N' headings found in the input file.", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
total_pages = len(pages)
|
||
sorted_pages = sorted(pages.keys())
|
||
print(f"Found {total_pages} pages (Page {sorted_pages[0]} – {sorted_pages[-1]}).")
|
||
|
||
# ── Load existing results (resume support) ────────────────────────────────
|
||
results: dict[str, dict] = {}
|
||
if output_path.exists():
|
||
try:
|
||
results = json.loads(output_path.read_text(encoding="utf-8"))
|
||
already_done = len(results)
|
||
print(f"Resuming — {already_done} page(s) already scored, skipping them.")
|
||
except (json.JSONDecodeError, OSError) as exc:
|
||
warnings.warn(f"Could not read existing output ({exc}); starting fresh.")
|
||
results = {}
|
||
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# ── Score each page ───────────────────────────────────────────────────────
|
||
scored_this_run = 0
|
||
for page_num in sorted_pages:
|
||
page_key = str(page_num)
|
||
|
||
if page_key in results:
|
||
continue # already scored — skip
|
||
|
||
page_text = pages[page_num]
|
||
result = score_page(page_text, page_num, api_key)
|
||
results[page_key] = result
|
||
scored_this_run += 1
|
||
|
||
# Progress line
|
||
score_str = str(result["score"]) if result["score"] > 0 else "ERR"
|
||
print(f" Page {page_num}/{sorted_pages[-1]} — score: {score_str}")
|
||
|
||
# Write after every page so a crash loses minimal work
|
||
output_path.write_text(
|
||
json.dumps(results, indent=2, ensure_ascii=False),
|
||
encoding="utf-8"
|
||
)
|
||
|
||
if scored_this_run > 0 and page_num != sorted_pages[-1]:
|
||
time.sleep(args.delay)
|
||
|
||
# ── Final save & summary ──────────────────────────────────────────────────
|
||
output_path.write_text(
|
||
json.dumps(results, indent=2, ensure_ascii=False),
|
||
encoding="utf-8"
|
||
)
|
||
print(f"\nResults written to: {output_path}")
|
||
print_summary(results)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|