app/recognizer.py

"""Chargement et inférence du modèle CRNN fine-tuné sur les écrans de pompe.

C'est le "modèle IA développé" au sens du §8.4 du cahier des charges :
un CRNN (crnn_vgg16_bn, doctr) fine-tuné sur nos propres données annotées
(voir train/)
"""
import os
from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np
import torch
import torch.nn.functional as F

from .preprocessing import FIELD_NAMES, split_lcd_lines

try:
    from doctr.datasets import VOCABS
    from doctr.models import crnn_vgg16_bn
except Exception as exc:  # pragma: no cover - dépend de l'environnement
    VOCABS = None
    crnn_vgg16_bn = None
    _DOCTR_IMPORT_ERROR = exc
else:
    _DOCTR_IMPORT_ERROR = None

VOCAB = VOCABS["french"] if VOCABS is not None else None
IMG_H = 32
MODULE_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_MODEL_PATH = os.path.join(MODULE_ROOT, "models", "crnn_fuel_pump_best.pt")


def get_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_model(model_path: Optional[str] = None, device: Optional[torch.device] = None):
    if model_path is None:
        model_path = DEFAULT_MODEL_PATH

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Modèle introuvable : {model_path}")

    if crnn_vgg16_bn is None or VOCAB is None:
        raise ImportError(
            "python-doctr n'est pas installé (voir requirements.txt)."
        ) from _DOCTR_IMPORT_ERROR

    if device is None:
        device = get_device()

    model = crnn_vgg16_bn(pretrained=False, pretrained_backbone=False, vocab=VOCAB)
    model.load_state_dict(torch.load(model_path, map_location=device))
    return model.to(device).eval()


def _resize_preserve_aspect(tensor: torch.Tensor, img_h: int) -> torch.Tensor:
    """Redimensionne en conservant le ratio d'aspect (hauteur fixe). Doit
    rester cohérent avec le prétraitement utilisé à l'entraînement
    (train/finetune_doctr.py::_resize_preserve_aspect) : un écart entre les
    deux dégrade silencieusement la précision du modèle.
    """
    _, h, w = tensor.shape
    new_w = max(16, int(img_h * w / max(h, 1)))
    resized = F.interpolate(
        tensor.unsqueeze(0), size=(img_h, new_w),
        mode="bilinear", align_corners=False,
    ).squeeze(0)
    mn, mx = resized.min(), resized.max()
    if mx - mn > 0.01:
        resized = (resized - mn) / (mx - mn)
    return resized.unsqueeze(0)


def preprocess_crop(crop_bgr: np.ndarray, img_h: int = IMG_H) -> torch.Tensor:
    img = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    tensor = torch.from_numpy(img).permute(2, 0, 1)
    return _resize_preserve_aspect(tensor, img_h)


@torch.no_grad()
def predict_line(model, crop_bgr: np.ndarray, device: torch.device) -> Tuple[str, float]:
    tensor = preprocess_crop(crop_bgr)
    out = model(tensor.to(device), return_preds=True)
    preds = out.get("preds", [])
    if not preds:
        return "", 0.0
    text, confidence = preds[0]
    return text, max(float(confidence), 0.0)


def recognize_screen(model, screen_crop: np.ndarray, device: torch.device,
                     n_lines: int = 3) -> List[Dict[str, Any]]:
    """Découpe le crop d'écran en lignes puis lit chaque valeur avec le CRNN.

    Retourne une liste de {'field', 'text', 'confidence'} — un élément par
    ligne détectée (prix / volume / prix_litre), dans l'ordre d'affichage.
    """
    line_crops, _ = split_lcd_lines(screen_crop, n_lines=n_lines)
    results = []
    for idx, crop in enumerate(line_crops):
        if crop.size == 0:
            continue
        text, confidence = predict_line(model, crop, device)
        field = FIELD_NAMES[idx] if idx < len(FIELD_NAMES) else f"ligne_{idx}"
        results.append({"field": field, "text": text, "confidence": confidence})
    return results