src/ocr/surya_engine.py

"""
Moteur OCR Surya — extraction texte ligne par ligne.
Usage CPU uniquement.
"""
import os
os.environ.setdefault("SURYA_TORCH_DEVICE", "cpu")
os.environ.setdefault("TORCH_DEVICE", "cpu")

import logging
import time
from PIL import Image

logger = logging.getLogger(__name__)

LANGS = ["fr", "en", "ar"]

_surya_det = None
_surya_rec = None


def get_surya_models():
    global _surya_det, _surya_rec
    if _surya_det is None or _surya_rec is None:
        logger.info("Chargement modèles Surya (CPU)…")
        t0 = time.time()
        from surya.models import DetectionPredictor
        from surya.recognition import RecognitionPredictor
        _surya_det = DetectionPredictor(device="cpu")
        _surya_rec = RecognitionPredictor(device="cpu")
        logger.info("Surya prêt en %.1fs", time.time() - t0)
    return _surya_det, _surya_rec


def run_surya_ocr(pil_image: Image.Image) -> dict:
    det, rec = get_surya_models()
    img = pil_image.convert("RGB")
    page_w, page_h = img.size

    t0 = time.time()
    rec_preds = rec([img], [LANGS], det_predictor=det)
    elapsed = time.time() - t0

    words = []
    for line in rec_preds[0].text_lines:
        text = line.text.strip()
        if not text:
            continue
        if hasattr(line, "bbox") and line.bbox:
            x1, y1, x2, y2 = line.bbox
        else:
            x1, y1, x2, y2 = 0, 0, page_w, page_h

        box = [
            int(x1 / page_w * 1000),
            int(y1 / page_h * 1000),
            int(x2 / page_w * 1000),
            int(y2 / page_h * 1000),
        ]
        words.append({
            "text":       text,
            "box":        box,
            "confidence": round(line.confidence, 3),
            "is_arabic":  False,
        })

    return {
        "words":     words,
        "page_w":    page_w,
        "page_h":    page_h,
        "full_text": "\n".join(w["text"] for w in words),
        "n_lines":   len(words),
        "ocr_time":  round(elapsed, 2),
    }