Create 03_ocr_bakeoff.py

smoke_signal/scripts/03_ocr_bakeoff.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+"""
+Smoke Signal - Stage 4: OCR Bake-Off
+Run Surya (primary) and Tesseract (fallback) on rendered pages.
+Freezes baseline OCR config after comparison.
+
+Usage:
+    python scripts/03_ocr_bakeoff.py
+    python scripts/03_ocr_bakeoff.py --book-id SS-BOOK-0001
+    python scripts/03_ocr_bakeoff.py --engine surya
+    python scripts/03_ocr_bakeoff.py --engine tesseract
+    python scripts/03_ocr_bakeoff.py --engine both
+    python scripts/03_ocr_bakeoff.py --dry-run
+"""
+
+import argparse
+import csv
+import json
+import sys
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+ROOT         = Path(__file__).resolve().parents[1]
+MANIFEST_CSV = ROOT / "manifest" / "source_manifest.csv"
+PROFILES_DIR = ROOT / "manifest" / "page_profiles"
+RENDERS_DIR  = ROOT / "renders"
+OCR_RAW_DIR  = ROOT / "ocr_raw"
+CONFIGS_DIR  = ROOT / "configs"
+LOGS_DIR     = ROOT / "logs"
+
+for d in [OCR_RAW_DIR, CONFIGS_DIR, LOGS_DIR]:
+    d.mkdir(parents=True, exist_ok=True)
+
+OCR_CONFIG = {
+    "config_version":      "ss_ocr_config_v0.1",
+    "primary_engine":      "surya",
+    "fallback_engine":     "tesseract",
+    "surya_langs":         ["en"],
+    "surya_det_batch":     4,
+    "surya_rec_batch":     4,
+    "tesseract_lang":      "eng",
+    "tesseract_psm":       6,
+    "confidence_threshold_auto_accept": 0.85,
+    "confidence_threshold_review":      0.60,
+    "confidence_threshold_quarantine":  0.40,
+    "eligible_routes":     ["ocr", "hybrid"],
+}
+
+
+def load_manifest():
+    records = {}
+    if not MANIFEST_CSV.exists():
+        return records
+    with open(MANIFEST_CSV, newline="", encoding="utf-8") as f:
+        for row in csv.DictReader(f):
+            if row.get("book_id"):
+                records[row["book_id"]] = row
+    return records
+
+
+def save_manifest(records):
+    fields = ["book_id","source_id","filename","sha256","file_size_bytes",
+              "page_count","rights_class","source_location","acquisition_date",
+              "status","allowed_use","notes"]
+    rows = sorted(records.values(), key=lambda r: r.get("book_id",""))
+    with open(MANIFEST_CSV, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(f, fieldnames=fields)
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def load_page_profile(book_id):
+    p = PROFILES_DIR / f"{book_id}_page_profile.json"
+    if not p.exists():
+        return None
+    with open(p, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def run_surya(image_path, langs):
+    try:
+        from PIL import Image
+        from surya.ocr import run_ocr
+        from surya.model.detection.model import load_model as load_det
+        from surya.model.detection.processor import load_processor as load_det_proc
+        from surya.model.recognition.model import load_model as load_rec
+        from surya.model.recognition.processor import load_processor as load_rec_proc
+    except ImportError as e:
+        return {"engine":"surya","error":str(e),"text":"","words":[],"confidence":0.0}
+    try:
+        img = Image.open(str(image_path)).convert("RGB")
+        det_m, det_p = load_det(), load_det_proc()
+        rec_m, rec_p = load_rec(), load_rec_proc()
+        results = run_ocr([img], [langs], det_m, det_p, rec_m, rec_p)
+        page = results[0]
+        words, confs, lines = [], [], []
+        for line in page.text_lines:
+            t = line.text.strip()
+            c = float(line.confidence) if hasattr(line,"confidence") else 1.0
+            if t:
+                lines.append(t)
+                confs.append(c)
+                words.append({"text":t,"confidence":round(c,4),"bbox":getattr(line,"bbox",None)})
+        avg = round(sum(confs)/len(confs),4) if confs else 0.0
+        return {"engine":"surya","text":"\n".join(lines),"words":words,"confidence":avg,"error":None}
+    except Exception as e:
+        return {"engine":"surya","error":str(e),"text":"","words":[],"confidence":0.0}
+
+
+def run_tesseract(image_path, lang, psm):
+    try:
+        import pytesseract
+        from PIL import Image
+    except ImportError as e:
+        return {"engine":"tesseract","error":str(e),"text":"","words":[],"confidence":0.0}
+    try:
+        img = Image.open(str(image_path)).convert("RGB")
+        cfg = f"--psm {psm}"
+        data = pytesseract.image_to_data(img, lang=lang, config=cfg, output_type=pytesseract.Output.DICT)
+        words, confs = [], []
+        for i, t in enumerate(data["text"]):
+            t = str(t).strip()
+            c = int(data["conf"][i])
+            if t and c > 0:
+                cn = c/100.0
+                words.append({"text":t,"confidence":round(cn,4),"bbox":[data["left"][i],data["top"][i],data["left"][i]+data["width"][i],data["top"][i]+data["height"][i]]})
+                confs.append(cn)
+        avg = round(sum(confs)/len(confs),4) if confs else 0.0
+        text = pytesseract.image_to_string(img, lang=lang, config=cfg).strip()
+        return {"engine":"tesseract","text":text,"words":words,"confidence":avg,"error":None}
+    except Exception as e:
+        return {"engine":"tesseract","error":str(e),"text":"","words":[],"confidence":0.0}
+
+
+def confidence_gate(conf, config):
+    if conf >= config["confidence_threshold_auto_accept"]: return "auto-accept"
+    if conf >= config["confidence_threshold_review"]:      return "review-required"
+    return "quarantine"
+
+
+def ocr_page(image_path, book_id, page_num, engine, config, dry_run=False):
+    result = {"book_id":book_id,"page_number":page_num,"image_path":str(image_path),
+              "engine":engine,"ocr_at":datetime.utcnow().isoformat()+"Z",
+              "config_version":config["config_version"]}
+    if dry_run:
+        result.update({"text":"[dry-run]","confidence":0.0,"gate":"dry-run","error":None,"words":[]})
+        return result
+    ocr_out = run_surya(image_path, config["surya_langs"]) if engine == "surya" else \
+              run_tesseract(image_path, config["tesseract_lang"], config["tesseract_psm"])
+    result.update(ocr_out)
+    result["gate"] = confidence_gate(result.get("confidence",0.0), config)
+    out_dir = OCR_RAW_DIR / book_id
+    out_dir.mkdir(parents=True, exist_ok=True)
+    with open(out_dir / f"{book_id}_page_{page_num:04d}_{engine}_ocr.json","w",encoding="utf-8") as f:
+        json.dump(result, f, indent=2, ensure_ascii=False)
+    return result
+
+
+def ocr_book(record, engine, dry_run=False):
+    book_id = record["book_id"]
+    print(f"\n  [{book_id}] {record['filename']} — {engine}")
+    profile = load_page_profile(book_id)
+    if not profile:
+        print("    No profile. Run 02_profile_pdfs.py first.")
+        return {"book_id":book_id,"error":"no_profile","pages":[]}
+    eligible = [p for p in profile["pages"] if p.get("route") in OCR_CONFIG["eligible_routes"]]
+    print(f"    OCR pages: {len(eligible)} / {profile['page_count']}")
+    results, confs, gates, errors = [], [], {"auto-accept":0,"review-required":0,"quarantine":0,"dry-run":0}, []
+    for p in eligible:
+        pn = p["page_number"]
+        rp = p.get("render_path")
+        if not rp:
+            cands = list((RENDERS_DIR/book_id).glob(f"{book_id}_page_{pn:04d}_*.png")) if (RENDERS_DIR/book_id).exists() else []
+            rp = str(cands[0]) if cands else None
+        if not rp:
+            errors.append({"page":pn,"error":"no_render"}); continue
+        ip = Path(rp) if Path(rp).is_absolute() else ROOT/rp
+        if not ip.exists():
+            errors.append({"page":pn,"error":"render_missing"}); continue
+        r = ocr_page(ip, book_id, pn, engine, OCR_CONFIG, dry_run)
+        results.append(r)
+        c = r.get("confidence",0.0)
+        g = r.get("gate","quarantine")
+        confs.append(c)
+        gates[g] = gates.get(g,0)+1
+        sym = "v" if g=="auto-accept" else "!" if g=="review-required" else "x"
+        print(f"    {sym} p{pn:03d} conf={c:.2f} gate={g}")
+    avg = round(sum(confs)/len(confs),4) if confs else 0.0
+    print(f"    Avg={avg:.2f} | {gates}")
+    return {"book_id":book_id,"engine":engine,"pages_ocred":len(results),"avg_confidence":avg,"gate_counts":gates,"errors":errors,"error":None}
+
+
+def save_config(config):
+    p = CONFIGS_DIR / f"{config['config_version']}.json"
+    if not p.exists():
+        with open(p,"w",encoding="utf-8") as f:
+            json.dump({**config,"frozen_at":datetime.utcnow().isoformat()+"Z"},f,indent=2)
+        print(f"  Config frozen -> {p.relative_to(ROOT)}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Smoke Signal Stage 4: OCR Bake-Off")
+    parser.add_argument("--book-id")
+    parser.add_argument("--batch-id")
+    parser.add_argument("--engine", choices=["surya","tesseract","both"], default="surya")
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--all", action="store_true")
+    args = parser.parse_args()
+    run_id  = args.batch_id or f"SS-RUN-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
+    engines = ["surya","tesseract"] if args.engine == "both" else [args.engine]
+    print(f"\nSmoke Signal Stage 4 | {run_id} | engines={engines}")
+    manifest = load_manifest()
+    if not manifest:
+        print("Manifest empty. Run stage 1 first."); sys.exit(1)
+    eligible = ["profiled","rendered"] if not args.all else ["profiled","rendered","ocred"]
+    books = [manifest[args.book_id]] if args.book_id else [r for r in manifest.values() if r.get("status") in eligible]
+    if not books:
+        print(f"No books in status {eligible}. Run stage 3 first."); sys.exit(0)
+    print(f"Books: {len(books)}")
+    all_results = []
+    t0 = time.time()
+    for rec in books:
+        for eng in engines:
+            res = ocr_book(rec, eng, args.dry_run)
+            all_results.append(res)
+            if not res.get("error") and not args.dry_run:
+                manifest[rec["book_id"]]["status"] = "ocred"
+    if not args.dry_run:
+        save_manifest(manifest)
+        save_config(OCR_CONFIG)
+        log = LOGS_DIR / f"{run_id}_ocr_bakeoff.json"
+        with open(log,"w") as f:
+            json.dump({"run_id":run_id,"engines":engines,"config":OCR_CONFIG,"results":all_results},f,indent=2)
+        print(f"  Log -> {log.relative_to(ROOT)}")
+    elapsed = round(time.time()-t0,1)
+    succeeded = sum(1 for r in all_results if not r.get("error"))
+    total_pages = sum(r.get("pages_ocred",0) for r in all_results)
+    avg = sum(r.get("avg_confidence",0) for r in all_results if not r.get("error"))/max(succeeded,1)
+    print(f"\nDone: {succeeded}/{len(all_results)} books | {total_pages} pages | avg_conf={avg:.2f} | {elapsed}s")
+    print("Next: run 04_region_detector.py (Stage 5)")
+
+
+if __name__ == "__main__":
+    main()