scan4llm/src/ocr/surya_engine.py

from __future__ import annotations

import functools
import logging
import os
import shutil
from pathlib import Path
from typing import TYPE_CHECKING

from src.env import load_env
from src.image_utils import calc_scale_dims
from src.models import OcrPageResult, ParsedBlock

if TYPE_CHECKING:
    from PIL import Image as PILImage

logger = logging.getLogger(__name__)

# CPU optimizations
_SURYA_CTX_PER_SLOT: int = 8192
_SURYA_PARALLEL: int = 1
_IMAGE_MAX_WIDTH: int = 1056


def _setup_surya_env() -> None:
    load_env()
    os.environ.setdefault("SURYA_INFERENCE_BACKEND", "llamacpp")
    os.environ.setdefault("SURYA_INFERENCE_KEEP_ALIVE", "true")
    os.environ.setdefault("SURYA_GUIDED_LAYOUT", "false")
    os.environ.setdefault("SURYA_INFERENCE_PARALLEL", str(_SURYA_PARALLEL))
    os.environ.setdefault("SURYA_INFERENCE_CTX_PER_SLOT", str(_SURYA_CTX_PER_SLOT))
    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
    cpu_count = os.cpu_count() or 4
    os.environ.setdefault(
        "LLAMA_CPP_EXTRA_ARGS", f"-t {cpu_count} --threads-batch {cpu_count}"
    )


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

    Optimizations for CPU:
      - llama-server kept alive between calls (SURYA_INFERENCE_KEEP_ALIVE=true)
      - reduced context per slot (8192 vs 12288) saves memory
      - single parallel slot (no batch needed for single page)
      - image pre-scaled to 1056px wide (≈96 DPI)
      - explicit thread count matching CPU cores
    """

    def __init__(self) -> None:
        _setup_surya_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)
        return self._process_image(image)

    def _health_check(self) -> bool:
        import requests

        port = os.environ.get("SURYA_INFERENCE_PORT", "")
        host = os.environ.get("SURYA_INFERENCE_HOST", "127.0.0.1")
        if not port:
            return True
        try:
            r = requests.get(f"http://{host}:{port}/health", timeout=10)
            return r.status_code == 200
        except requests.RequestException:
            logger.warning("llama-server health check failed")
            return False

    def _process_image(self, image: PILImage.Image) -> OcrPageResult:
        scaled = _pre_scale(image)
        results = self._predictor([scaled])
        if results:
            return self._to_result(results[0], scaled)
        return OcrPageResult()

    @staticmethod
    def _to_result(page, image: PILImage.Image) -> OcrPageResult:
        blocks: list[ParsedBlock] = []
        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={},
            width=image.width,
            height=image.height,
        )


def _pre_scale(image: PILImage.Image) -> PILImage.Image:
    new_w, new_h = calc_scale_dims(*image.size, _IMAGE_MAX_WIDTH)
    if (new_w, new_h) != image.size:
        return image.resize((new_w, new_h), 1)  # PIL.Image.LANCZOS
    return image


@functools.lru_cache(maxsize=1)
def _get_surya_engine() -> SuryaEngine:
    return SuryaEngine()


def surya_ocr_image(image_path: str) -> OcrPageResult:
    engine = _get_surya_engine()
    if not engine._health_check():
        logger.warning("llama-server не отвечает, перезапуск...")
        _get_surya_engine.cache_clear()
        engine = _get_surya_engine()
        if not engine._health_check():
            raise RuntimeError(
                "llama-server недоступен после перезапуска. "
                "Проверьте: `ps aux | grep llama-server`. "
                "Попробуйте: `llama-server --version`"
            )
    return engine.process(image_path)