ocr.py

import io
import os
import re

import numpy as np
from PIL import Image

# Facteur d'agrandissement appliqué aux crops de bulles AVANT l'OCR. Les
# bulles croppées sont souvent petites ; agrandir l'image (LANCZOS) améliore
# nettement la reconnaissance de PaddleOCR (et de manga-ocr) sur les polices
# manga stylisées et le texte de petite taille. Surchargeable via
# OCR_UPSCALE (1 = désactivé).
UPSCALE_FACTOR = int(os.environ.get("OCR_UPSCALE", "2"))


def _prepare_image(image_bytes: bytes) -> Image.Image:
    """Charge l'image et l'agrandit (LANCZOS) pour améliorer la précision OCR."""
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    if UPSCALE_FACTOR > 1:
        w, h = image.size
        image = image.resize(
            (w * UPSCALE_FACTOR, h * UPSCALE_FACTOR), Image.LANCZOS
        )
    return image


def _clean_ocr_text(text: str) -> str:
    """Nettoyage LÉGER du texte OCR.

    Ne corrige PAS les mots (risque de casser des noms propres ou des mots
    composés), mais retire le bruit évident : espaces multiples, espaces
    avant ponctuation, espaces en début/fin.
    """
    text = re.sub(r"\s+", " ", text)
    text = re.sub(r"\s+([,.;:!?…])", r"\1", text)
    return text.strip()


# Mapping des codes ISO 639-1 (apps/web/lib/languages.ts) vers les codes de
# langue de PaddleOCR 2.x. Les langues à alphabet latin sans modèle dédié
# (espagnol, italien, portugais, néerlandais...) utilisent le modèle
# générique `latin`. Le japonais (ja) utilise le modèle dédié `japan`.
PADDLE_LANG_MAP = {
    "ja": "japan",
    "de": "german",
    "en": "en",
    "es": "latin",
    "it": "latin",
    "pt": "latin",
    "nl": "latin",
    "fr": "french",
    "ru": "cyrillic",
    "ko": "korean",
    "zh": "ch",
}


class PaddleOcr:
    """
    OCR (PaddleOCR) pour TOUTES les langues, dont le japonais (lang='japan'),
    exécuté côté serveur. Bien plus robuste que Tesseract.js sur les polices
    manga stylisées, et déporté côté serveur — aligné avec l'objectif de mise
    à disposition publique (l'utilisateur final n'a rien à installer).

    NB : manga-ocr (plus précis sur le japonais) a été retiré car il dépend
    de transformers, incompatible avec gradio 6.24.0 imposé par HF Spaces
    ZeroGPU (conflit huggingface-hub). Le japonais passe donc par PaddleOCR.

    Les modèles PaddleOCR se téléchargent automatiquement depuis les serveurs
    de PaddleOCR au premier appel, puis restent en cache dans le conteneur.
    """

    def __init__(self) -> None:
        self._reader = None
        self._lang = None

    def _ensure_loaded(self, lang: str):
        # Un lecteur PaddleOCR est lié à une langue : on en recrée un si la
        # langue demandée change (coût unique, mis en cache ensuite).
        if self._reader is None or self._lang != lang:
            from paddleocr import PaddleOCR

            self._reader = PaddleOCR(
                use_angle_cls=True,
                lang=lang,
                show_log=False,
            )
            self._lang = lang
        return self._reader

    def read(self, image_bytes: bytes, lang: str) -> str:
        reader = self._ensure_loaded(lang)
        image = _prepare_image(image_bytes)
        result = reader.ocr(np.array(image), cls=True)

        # PaddleOCR 2.x renvoie une liste de listes : pour une seule image,
        # result[0] est la liste des détections, chacune au format
        # [box, (texte, confiance)]. On concatène les textes détectés.
        lines: list[str] = []
        if result and result[0]:
            for detection in result[0]:
                text = detection[1][0]
                if text and text.strip():
                    lines.append(_clean_ocr_text(text))
        return "\n".join(lines)