Book-Reader-Assist/tools/ocr_image_pages.py

"""OCR PDF pages that have no native extractable text.

Usage: python tools/ocr_image_pages.py "<Book Title>"

Writes library/<Book Title>/ocr-image-pages.json. The split_book.py exporter
automatically inserts this OCR beneath the matching PDF page markers.
"""
import json
import sys
from pathlib import Path

import numpy as np
import pymupdf
from rapidocr_onnxruntime import RapidOCR


ROOT = Path(__file__).resolve().parent.parent / "library"


def main(book):
    pdf_path = ROOT / book / "source-file.pdf"
    output_path = ROOT / book / "ocr-image-pages.json"
    document = pymupdf.open(pdf_path)
    pages = [number + 1 for number, page in enumerate(document) if not page.get_text().strip()]
    engine = RapidOCR()
    recovered = {}

    print(f"OCR pages for {book}: {pages}", flush=True)
    for index, page_number in enumerate(pages, start=1):
        page = document[page_number - 1]
        pixmap = page.get_pixmap(
            matrix=pymupdf.Matrix(2, 2),
            colorspace=pymupdf.csRGB,
            alpha=False,
        )
        image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(
            pixmap.height, pixmap.width, pixmap.n
        )
        result, _ = engine(image)
        lines = [item[1].strip() for item in (result or []) if item[1].strip()]
        recovered[str(page_number)] = (
            "\n".join(lines)
            if lines
            else "[Image-only page; OCR found no readable text. Consult the PDF for visual content.]"
        )
        print(f"{index}/{len(pages)} PDF page {page_number}: {len(lines)} OCR lines", flush=True)

    output_path.write_text(
        json.dumps(recovered, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
        newline="\n",
    )
    print(f"Wrote {output_path}", flush=True)


if __name__ == "__main__":
    main(sys.argv[1])