scan4llm/src/ocr/surya_engine.py

from __future__ import annotations

import logging
import os
import shutil
from pathlib import Path

from src.ocr.paddle_engine import OcrPageResult, ParsedBlock

logger = logging.getLogger(__name__)

_ENV_LOADED: bool = False


def _load_env() -> None:
    global _ENV_LOADED
    if _ENV_LOADED:
        return
    env_path = Path(__file__).resolve().parent.parent.parent / ".env"
    if env_path.exists():
        with open(env_path) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, _, value = line.partition("=")
                    key = key.strip()
                    if key not in os.environ:
                        os.environ[key] = value.strip().strip("\"'")
    _ENV_LOADED = True


def _ensure_env() -> None:
    _load_env()
    os.environ.setdefault("SURYA_INFERENCE_BACKEND", "llamacpp")
    os.environ.setdefault("SURYA_GUIDED_LAYOUT", "false")
    if "LLAMA_CPP_BINARY" not in os.environ:
        binary = shutil.which("llama-server") or os.path.expanduser(
            "~/.local/bin/llama-server"
        )
        if Path(binary).exists():
            os.environ["LLAMA_CPP_BINARY"] = binary


class SuryaEngine:
    """OCR engine using Surya 2 VLM via llama.cpp."""

    def __init__(self) -> None:
        _ensure_env()
        try:
            from surya.inference import SuryaInferenceManager
            from surya.recognition import RecognitionPredictor

            self._manager = SuryaInferenceManager()
            self._predictor = RecognitionPredictor(self._manager)
        except ImportError as e:
            msg = (
                "Surya engine not available. Install: pip install surya-ocr torch"
            )
            raise ImportError(msg) from e

    def process(self, image_path: str) -> OcrPageResult:
        from PIL import Image

        image = Image.open(image_path)
        results = self._predictor([image])

        blocks: list[ParsedBlock] = []
        raw_json: dict = {}

        if results:
            page = results[0]
            for blk in getattr(page, "blocks", []):
                html = getattr(blk, "html", "") or ""
                label = getattr(blk, "label", "text")
                bbox = getattr(blk, "bbox", [0, 0, 0, 0])
                confidence = float(getattr(blk, "confidence", 0.9))
                blocks.append(
                    ParsedBlock(
                        label=label,
                        content=html,
                        bbox=(int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])),
                        confidence=confidence,
                    )
                )

        return OcrPageResult(
            blocks=blocks,
            raw_json=raw_json,
            width=image.width,
            height=image.height,
        )


_surya_engine: SuryaEngine | None = None


def surya_ocr_image(image_path: str) -> OcrPageResult:
    global _surya_engine
    if _surya_engine is None:
        _surya_engine = SuryaEngine()
    return _surya_engine.process(image_path)