OCR-Viewer (#1)
Co-authored-by: nathan <nathan.kehler@gmail.com> Reviewed-on: #1
This commit is contained in:
298
scripts/ocr_confidence.py
Normal file
298
scripts/ocr_confidence.py
Normal file
@@ -0,0 +1,298 @@
|
||||
"""
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user