scripts/ocr_book.py

"""
SUPERSEDED by scripts/build_index_from_crawler.py — kept for reference only.
See README.md for the current build path.

OCR a scanned (image-only) book PDF into the same crawl-style text layout the
rest of the corpus uses, then update its state.json so build_index picks it up.

Renders each page to a grayscale image with PyMuPDF and runs the Tesseract CLI
(image via stdin -> text via stdout). Used for books that extract_books.py
flagged status="low_yield_scanned".

Usage:
    python ocr_book.py --pdf "../books/613142275-Understanding-Digital-Marketing.pdf"
    python ocr_book.py --pdf ... --out ../crawled_books --dpi 300 --lang eng
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path

import fitz  # PyMuPDF


def slugify(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")[:80] or "book"


def name_hash(name: str) -> str:
    return hashlib.sha1(name.encode("utf-8")).hexdigest()[:16]


def clean_text(text: str) -> str:
    text = text.replace("\x0c", "\n")
    text = re.sub(r"[ \t]+\n", "\n", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"(\w)-\n(\w)", r"\1\2", text)
    return text.strip()


def ocr_page(png_bytes: bytes, tesseract: str, lang: str) -> str:
    proc = subprocess.run(
        [tesseract, "stdin", "stdout", "-l", lang, "--psm", "1"],
        input=png_bytes,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
    )
    return proc.stdout.decode("utf-8", errors="ignore")


def main():
    parser = argparse.ArgumentParser(description="OCR a scanned book PDF -> crawl-style text")
    parser.add_argument("--pdf", required=True)
    parser.add_argument("--out", default="../crawled_books")
    parser.add_argument("--dpi", type=int, default=300)
    parser.add_argument("--lang", default="eng")
    parser.add_argument("--tesseract", default=shutil.which("tesseract") or "tesseract")
    args = parser.parse_args()

    pdf = Path(args.pdf)
    if not pdf.exists():
        sys.exit(f"PDF not found: {pdf}")

    slug = slugify(pdf.stem)
    h = name_hash(pdf.name)
    out_dir = Path(args.out) / slug
    pages_dir = out_dir / "pages"
    pages_dir.mkdir(parents=True, exist_ok=True)

    doc = fitz.open(str(pdf))
    n = doc.page_count
    zoom = args.dpi / 72.0
    mat = fitz.Matrix(zoom, zoom)

    print(f"OCR {pdf.name}: {n} pages @ {args.dpi} DPI (tesseract: {args.tesseract})", flush=True)
    parts = []
    t0 = time.time()
    for i in range(n):
        page = doc[i]
        pix = page.get_pixmap(matrix=mat, colorspace=fitz.csGRAY)
        png = pix.tobytes("png")
        txt = ocr_page(png, args.tesseract, args.lang)
        parts.append(txt)
        if (i + 1) % 20 == 0 or i + 1 == n:
            elapsed = time.time() - t0
            rate = (i + 1) / elapsed
            eta = (n - i - 1) / rate if rate else 0
            print(f"  {i+1:>4}/{n} pages  ({rate:.2f} pg/s, ETA {eta/60:.1f} min)", flush=True)
    doc.close()

    text = clean_text("\n".join(parts))
    (pages_dir / f"{h}.txt").write_text(text, encoding="utf-8")

    status = "ok_ocr" if len(text) >= 800 else "low_yield_ocr"
    state = {
        h: {
            "file": pdf.name,
            "category": "general",
            "status": status,
            "pages": n,
            "chars": len(text),
            "engine": f"tesseract-ocr@{args.dpi}dpi",
            "extracted_at": time.time(),
        }
    }
    (out_dir / "state.json").write_text(json.dumps(state, indent=2), encoding="utf-8")
    print(f"\nDone: {len(text):,} chars, status={status} -> {pages_dir / (h + '.txt')}", flush=True)


if __name__ == "__main__":
    main()