Spaces:
Sleeping
Sleeping
File size: 3,895 Bytes
b510add | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """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
|