app/ocr_worker.py

"""Standalone Surya OCR worker (runs in the isolated .venv-ocr).

Surya 0.22 pulls transformers 5.x + a llama.cpp backend, which conflicts with
the embedding stack (transformers 4.x). To keep those concerns decoupled, this
worker runs under its OWN interpreter (backend/.venv-ocr/bin/python) and is
invoked as a subprocess by app.ocr. It must NOT import the app package.

Usage:  python ocr_worker.py <input_file> <output_json>
Writes: {"lines": [{"text": str, "confidence": float}], "raw_text": str}
"""
import json
import re
import sys

from PIL import Image


def _load_images(path: str):
    if path.lower().endswith(".pdf"):
        import pypdfium2 as pdfium

        pdf = pdfium.PdfDocument(path)
        return [pdf[i].render(scale=2.0).to_pil().convert("RGB") for i in range(len(pdf))]
    return [Image.open(path).convert("RGB")]


def _strip_html(s: str | None) -> str:
    return re.sub(r"<[^>]+>", " ", s or "").strip()


def main() -> None:
    in_path, out_path = sys.argv[1], sys.argv[2]
    from surya.recognition import RecognitionPredictor

    rec = RecognitionPredictor()
    pages = rec(_load_images(in_path), full_page=True)

    lines = []
    for page in pages:
        for block in page.blocks:
            text = _strip_html(getattr(block, "html", None))
            if text:
                conf = getattr(block, "confidence", 1.0)
                lines.append({"text": text, "confidence": float(conf if conf is not None else 1.0)})

    with open(out_path, "w", encoding="utf-8") as f:
        json.dump({"lines": lines, "raw_text": "\n".join(l["text"] for l in lines)}, f)


if __name__ == "__main__":
    main()