inference.py

import io

from PIL import Image

MODEL_ID = "ogkalu/comic-text-and-bubble-detector"
CONFIDENCE_THRESHOLD = 0.3

# Classes du modèle (voir sa fiche HuggingFace) :
#   0 = bubble       (contour de bulle vide, sans texte à l'intérieur)
#   1 = text_bubble  (texte à l'intérieur d'une bulle)
#   2 = text_free    (texte hors bulle : onomatopées, cartouches, etc.)
# On ne garde que les zones contenant du texte : c'est ce qui doit devenir
# cliquable pour l'utilisateur. Les bulles vides (classe 0) sont ignorées.
RELEVANT_LABELS = {"text_bubble", "text_free"}


class BubbleDetector:
    """
    Détecte les zones de texte d'une page de manga avec le modèle
    open-source `ogkalu/comic-text-and-bubble-detector` (RT-DETR-v2, licence
    Apache 2.0), entraîné sur manga/webtoon/manhua/comics occidentaux.

    Utilise le pipeline `object-detection` standard de la librairie
    `transformers` : le pré-traitement (redimensionnement) et le
    post-traitement (décodage des boîtes, seuil de confiance) sont gérés en
    interne par HuggingFace, aucune logique de décodage écrite à la main ici
    — contrairement à des modèles YOLO bruts, le contrat d'API est stable et
    documenté.

    Le modèle (~170 Mo) se télécharge automatiquement depuis HuggingFace au
    premier appel, puis reste en cache — comme pour manga-ocr, aucune étape
    manuelle, aucune clé API.
    """

    def __init__(self) -> None:
        self._pipeline = None

    def _ensure_loaded(self):
        if self._pipeline is None:
            from transformers import pipeline
            import torch

            # GPU si disponible (RTX 4060), sinon CPU. device=0 = premier GPU.
            device = 0 if torch.cuda.is_available() else -1
            self._pipeline = pipeline(
                "object-detection", model=MODEL_ID, device=device
            )
        return self._pipeline

    def detect(self, image_bytes: bytes) -> list[dict]:
        detector = self._ensure_loaded()
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        predictions = detector(image, threshold=CONFIDENCE_THRESHOLD)

        boxes: list[dict] = []
        for pred in predictions:
            if pred["label"] not in RELEVANT_LABELS:
                continue
            box = pred["box"]  # {xmin, ymin, xmax, ymax} déjà en pixels réels
            boxes.append(
                {
                    "x": int(box["xmin"]),
                    "y": int(box["ymin"]),
                    "width": int(box["xmax"] - box["xmin"]),
                    "height": int(box["ymax"] - box["ymin"]),
                }
            )

        return boxes