""" Moteur OCR docTR — Optimisé CPU 8 Go RAM, Multilingue (FR / EN / AR). Modèles choisis pour CPU sans GPU : Détection : db_mobilenet_v3_large (~12 Mo, ~1-2 s/page) Reconnaissance : crnn_mobilenet_v3_large (~8 Mo, ~2-3 s/page) → Consommation mémoire totale : ≈ 350-450 Mo RAM → Temps traitement page A4 scannée : 3-6 s sur CPU standard Alternatives si encore trop lent : reco_arch = "crnn_mobilenet_v3_small" → ≈ 180 Mo, ~2 s/page Sortie : [{"text": str, "box": [x1,y1,x2,y2], "confidence": float, "is_arabic": bool}] Coordonnées normalisées 0-1000 (format LayoutLM / LayoutXLM). """ import gc import logging import os from typing import Union, Optional import numpy as np from PIL import Image, ImageDraw, ImageFont logger = logging.getLogger(__name__) # ───────────────────────────────────────────────────────────────────────────── # Constantes d'optimisation CPU # ───────────────────────────────────────────────────────────────────────────── # Largeur max envoyée à docTR pour l'inférence. # Au-delà, l'image est redimensionnée → gain de vitesse important. # 1024 px = bon compromis vitesse/précision sur CPU. MAX_INFERENCE_WIDTH = 512 # Architectures recommandées pour CPU 8 Go _DEFAULT_DET = "db_mobilenet_v3_large" # léger, rapide _DEFAULT_RECO = "crnn_mobilenet_v3_large" # bon équilibre précision/vitesse # Singleton : le prédicteur est chargé une seule fois par session Streamlit _predictor = None # ───────────────────────────────────────────────────────────────────────────── # Initialisation du prédicteur (singleton) # ───────────────────────────────────────────────────────────────────────────── def get_predictor( det_arch: str = _DEFAULT_DET, reco_arch: str = _DEFAULT_RECO, assume_straight_pages: bool = True, ): """ Charge et met en cache le prédicteur docTR (une seule fois par session). Optimisations CPU appliquées automatiquement : - Tous les threads disponibles (torch.set_num_threads) - Mode inférence uniquement (pas de gradient calculé) - Modèles légers MobileNetV3 Args: det_arch : Architecture de détection. 'db_mobilenet_v3_large' ← recommandé CPU (12 Mo) 'db_resnet50' ← précis mais lent sur CPU reco_arch : Architecture de reconnaissance. 'crnn_mobilenet_v3_large' ← recommandé CPU (8 Mo) 'crnn_mobilenet_v3_small' ← ultra-léger (3 Mo) 'crnn_vgg16_bn' ← moyen (60 Mo) 'parseq' ← précis mais lent (90 Mo) assume_straight_pages : True pour pages droites (plus rapide). Returns: doctr.models.OCRPredictor prêt à l'emploi. """ global _predictor if _predictor is not None: return _predictor # ── Import conditionnel (évite d'importer torch au démarrage) ───────────── try: import torch from doctr.models import ocr_predictor except ImportError: raise ImportError( "docTR n'est pas installé.\n" "Installez avec : pip install python-doctr[torch]\n" "Documentation : https://mindee.github.io/doctr/" ) # ── Libération mémoire avant chargement ────────────────────────────────── gc.collect() # ── Optimisations CPU ───────────────────────────────────────────────────── n_threads = min(os.cpu_count() or 2, 4) # plafonné à 4 pour limiter la RAM torch.set_num_threads(n_threads) torch.set_num_interop_threads(max(1, n_threads // 2)) logger.info( "Chargement docTR — det=%s reco=%s threads=%d", det_arch, reco_arch, n_threads, ) _predictor = ocr_predictor( det_arch=det_arch, reco_arch=reco_arch, pretrained=True, assume_straight_pages=assume_straight_pages, ) # Passer en mode évaluation (désactive dropout, batchnorm en mode train, etc.) _predictor.det_predictor.model.eval() _predictor.reco_predictor.model.eval() logger.info("Modèles docTR prêts (RAM ≈ 350-450 Mo).") return _predictor def reset_predictor(): """Libère les modèles de la mémoire (utile si on change d'architecture).""" global _predictor _predictor = None gc.collect() def trim_memory(): """ Demande au système d'exploitation de récupérer les pages mémoire inutilisées par PyTorch après une inférence CPU. Sans effet sur GPU. Sur Windows : SetProcessWorkingSetSize(-1, -1) vide le working set. Sur Linux : malloc_trim(0) via libc. """ gc.collect() gc.collect() try: import sys if sys.platform == "win32": import ctypes ctypes.windll.kernel32.SetProcessWorkingSetSize( ctypes.windll.kernel32.GetCurrentProcess(), -1, -1 ) else: ctypes.CDLL("libc.so.6").malloc_trim(0) except Exception: pass # ───────────────────────────────────────────────────────────────────────────── # Fonction principale d'OCR # ───────────────────────────────────────────────────────────────────────────── def run_ocr( image_input: Union[np.ndarray, Image.Image], det_arch: str = _DEFAULT_DET, reco_arch: str = _DEFAULT_RECO, assume_straight_pages: bool = True, min_confidence: float = 0.30, ) -> dict: """ Exécute l'OCR docTR sur une image prétraitée. L'image est redimensionnée à MAX_INFERENCE_WIDTH avant inférence pour réduire la charge mémoire et accélérer le traitement sur CPU. Args: image_input : numpy array (BGR ou gris) OU PIL.Image. det_arch : Architecture de détection. reco_arch : Architecture de reconnaissance. assume_straight_pages : Optimisation pages droites. min_confidence : Mots en dessous de ce seuil sont ignorés. Returns: dict : "words" → liste LayoutXLM (voir to_layoutxlm_format) "full_text" → texte complet (str) "page_w/h" → dimensions originales de la page "raw_result" → résultat brut docTR """ import io as _io try: import torch ctx = torch.inference_mode() # pas de calcul de gradient → économie RAM except ImportError: ctx = _NullContext() # ── Préparation image ───────────────────────────────────────────────────── pil_rgb = _to_rgb_pil(image_input) page_w, page_h = pil_rgb.size # Redimensionner pour l'inférence (gain de vitesse × 2-4 selon la taille) pil_inf = _resize_for_inference(pil_rgb, max_width=MAX_INFERENCE_WIDTH) # ── Inférence ───────────────────────────────────────────────────────────── predictor = get_predictor(det_arch, reco_arch, assume_straight_pages) try: from doctr.io import DocumentFile buf = _io.BytesIO() pil_inf.save(buf, format="PNG") buf.seek(0) doc = DocumentFile.from_images([buf.read()]) except Exception: doc = [np.array(pil_inf)] with ctx: raw_result = predictor(doc) # Libérer immédiatement le document (tenseurs d'entrée) del doc gc.collect() # ── Extraction ──────────────────────────────────────────────────────────── words = to_layoutxlm_format(raw_result, min_confidence=min_confidence) full_text = _reconstruct_text(raw_result) # Ne pas faire del raw_result ici : Python 3.12 lève UnboundLocalError si # predictor() a échoué auparavant. raw_result est libéré automatiquement # par le GC au retour de la fonction (plus aucune référence externe). del pil_inf gc.collect() gc.collect() # double appel : libère les cycles du premier passage return { "words": words, "full_text": full_text, "page_w": page_w, "page_h": page_h, } # ───────────────────────────────────────────────────────────────────────────── # Conversion vers format LayoutXLM (0-1000) # ───────────────────────────────────────────────────────────────────────────── def to_layoutxlm_format( doctr_result, min_confidence: float = 0.30, ) -> list: """ Convertit la sortie docTR en liste de dicts compatibles LayoutXLM. Format de sortie : [ { "text": "Facture", "box": [42, 80, 198, 112], # 0-1000 "confidence": 0.97, "is_arabic": False, }, ... ] Normalisation 0-1000 (standard LayoutLM v1/v2/XLM) : Les coordonnées relatives [0,1] de docTR sont multipliées par 1000. Args: doctr_result : Objet Document retourné par le prédicteur docTR. min_confidence : Seuil de confiance minimal. Returns: Liste de dicts triés de haut en bas puis gauche à droite. """ words_out = [] for page in doctr_result.pages: for block in page.blocks: for line in block.lines: for word in line.words: if word.confidence < min_confidence: continue geo = word.geometry # Extraire les coins (format bbox 2 pts ou polygone 4 pts) if len(geo) == 2: (rx1, ry1), (rx2, ry2) = geo else: xs = [p[0] for p in geo] ys = [p[1] for p in geo] rx1, ry1, rx2, ry2 = min(xs), min(ys), max(xs), max(ys) # Normalisation 0-1000 avec clamp de sécurité x1 = max(0, min(1000, int(rx1 * 1000))) y1 = max(0, min(1000, int(ry1 * 1000))) x2 = max(0, min(1000, int(rx2 * 1000))) y2 = max(0, min(1000, int(ry2 * 1000))) text = word.value.strip() if not text: continue words_out.append({ "text": text, "box": [x1, y1, x2, y2], "confidence": round(word.confidence, 4), "is_arabic": _is_arabic(text), }) # Tri : lecture naturelle haut → bas, gauche → droite words_out.sort(key=lambda w: (w["box"][1], w["box"][0])) return words_out # ───────────────────────────────────────────────────────────────────────────── # Fusion de mots en blocs de lignes (Grouping pour LayoutXLM) # ───────────────────────────────────────────────────────────────────────────── def merge_words_into_lines( words: list, y_thresh: int = 15, x_gap_thresh: int = 40, ) -> list: """ Regroupe les mots qui appartiennent à la même ligne visuelle en blocs fusionnés. Pourquoi ? docTR renvoie les mots un par un (ex : une désignation → 15 tokens isolés). LayoutXLM fonctionne mieux sur des blocs de ligne cohérents car : - Chaque bloc consomme 1 slot sur les 512 disponibles (au lieu de 15). - Le contexte spatial est plus lisible pour le modèle. Algorithme (3 étapes) : 1. Tri par Y_centre puis X → ordre de lecture naturel. 2. Regroupement vertical : deux mots forment la même ligne si |Y_centre_A − Y_centre_B| ≤ y_thresh. 3. Fusion horizontale : deux mots adjacents sur la même ligne sont fusionnés si le gap horizontal (x1_droit − x2_gauche) ≤ x_gap_thresh. Au-delà du seuil, une nouvelle colonne commence. Paramètres : y_thresh (int, 0-1000) : Tolérance verticale en unités normalisées. Recommandé : 10-20 (1-2 % de la hauteur). • 15 → absorbe la légère inclinaison / skew. • Trop grand (>30) → deux lignes adjacentes pourraient être fusionnées par erreur. x_gap_thresh (int, 0-1000) : Distance horizontale max entre deux mots pour les fusionner. Recommandé : 30-60 (3-6 % de la largeur). • 40 → fusionne les mots d'une même phrase. • Garde les gaps inter-colonnes (>60) séparés. Args: words : Liste [{"text": str, "box": [x1,y1,x2,y2], "confidence": float, "is_arabic": bool}] — sortie de to_layoutxlm_format(). Returns: Liste de blocs fusionnés, même format que words, triée en ordre de lecture. """ if not words: return words # ── Étape 1 : tri initial par Y_centre puis X (ordre lecture naturel) ───── sorted_words = sorted( words, key=lambda w: ((w["box"][1] + w["box"][3]) / 2, w["box"][0]), ) # ── Étape 2 : regroupement vertical en lignes ───────────────────────────── # Chaque élément de `lines` est une liste de mots sur la même ligne. lines: list[list[dict]] = [] for word in sorted_words: y_c = (word["box"][1] + word["box"][3]) / 2 placed = False # Chercher une ligne existante compatible (parcours en sens inverse # pour comparer à la ligne la plus récente en priorité) for line in reversed(lines): line_y_c = sum((w["box"][1] + w["box"][3]) / 2 for w in line) / len(line) if abs(y_c - line_y_c) <= y_thresh: line.append(word) placed = True break if not placed: lines.append([word]) # ── Étape 3 : tri gauche→droite dans chaque ligne ───────────────────────── for line in lines: line.sort(key=lambda w: w["box"][0]) # ── Étape 4 : fusion horizontale par ligne ──────────────────────────────── merged: list[dict] = [] for line in lines: group: dict | None = None for word in line: if group is None: group = { "text": word["text"], "box": list(word["box"]), # copie mutable "confidence": word["confidence"], "is_arabic": word.get("is_arabic", False), } else: # Gap horizontal : distance entre la fin du groupe et le début du mot gap = word["box"][0] - group["box"][2] if gap <= x_gap_thresh: # ── Fusion : concaténer le texte et étendre la bbox ─────── group["text"] += " " + word["text"] # Bbox englobante : min des coins haut-gauche, max des coins bas-droit group["box"][0] = min(group["box"][0], word["box"][0]) group["box"][1] = min(group["box"][1], word["box"][1]) group["box"][2] = max(group["box"][2], word["box"][2]) group["box"][3] = max(group["box"][3], word["box"][3]) # Confiance = pire confiance du groupe (prudence) group["confidence"] = min(group["confidence"], word["confidence"]) group["is_arabic"] = group["is_arabic"] or word.get("is_arabic", False) else: # ── Nouveau bloc (colonne différente ou espacement trop grand) merged.append(group) group = { "text": word["text"], "box": list(word["box"]), "confidence": word["confidence"], "is_arabic": word.get("is_arabic", False), } if group is not None: merged.append(group) # ── Étape 5 : tri final en ordre de lecture (haut→bas, gauche→droite) ───── merged.sort(key=lambda b: (b["box"][1], b["box"][0])) logger.debug( "merge_words_into_lines : %d mots → %d blocs (y_thresh=%d, x_gap=%d)", len(words), len(merged), y_thresh, x_gap_thresh, ) return merged # ───────────────────────────────────────────────────────────────────────────── # Visualisation — bounding boxes annotées # ───────────────────────────────────────────────────────────────────────────── def draw_bounding_boxes( image_input: Union[np.ndarray, Image.Image], words: list, color_latin: str = "#0088DD", color_arabic: str = "#E05A00", line_width: int = 2, show_text: bool = True, show_confidence: bool = False, ) -> Image.Image: """ Dessine les bounding boxes sur l'image pour vérification visuelle. Couleurs : Bleu (#0088DD) → texte latin (français / anglais) Orange (#E05A00) → texte arabe (détecté automatiquement) Args: image_input : Image source (PIL ou numpy BGR/gris). words : Sortie de to_layoutxlm_format(). color_latin : Couleur bbox texte latin. color_arabic : Couleur bbox texte arabe. line_width : Épaisseur contour (pixels). show_text : Afficher le texte reconnu au-dessus de la bbox. show_confidence : Ajouter le score de confiance dans le label. Returns: PIL.Image RGB annotée. """ pil = _to_rgb_pil(image_input).copy() draw = ImageDraw.Draw(pil) w_img, h_img = pil.size font_size = max(11, w_img // 90) try: font = ImageFont.truetype("arial.ttf", size=font_size) except IOError: font = ImageFont.load_default() for word in words: x1, y1, x2, y2 = word["box"] # Dénormalisation 0-1000 → pixels px1 = int(x1 * w_img / 1000) py1 = int(y1 * h_img / 1000) px2 = int(x2 * w_img / 1000) py2 = int(y2 * h_img / 1000) color = color_arabic if word.get("is_arabic") else color_latin draw.rectangle([px1, py1, px2, py2], outline=color, width=line_width) if show_text: label = word["text"] if show_confidence: label += f" {word['confidence']:.2f}" try: bbox_t = draw.textbbox((px1, py1), label, font=font) draw.rectangle( [bbox_t[0]-2, bbox_t[1]-2, bbox_t[2]+2, bbox_t[3]+2], fill=color, ) draw.text((px1, py1), label, fill="white", font=font, anchor="lb") except Exception: pass # Si la police ne supporte pas le glyphe arabe, on ignore return pil # ───────────────────────────────────────────────────────────────────────────── # Pipeline complet (prétraitement → OCR → LayoutXLM → visualisation) # ───────────────────────────────────────────────────────────────────────────── def ocr_pipeline( file_input, language_hint: Optional[str] = None, min_confidence: float = 0.30, apply_cleaning: bool = True, ) -> dict: """ Pipeline tout-en-un optimisé CPU 8 Go : 1. Prétraitement (image_preprocessor) 2. OCR léger docTR (MobileNetV3) 3. Formatage LayoutXLM 4. Visualisation bounding boxes 5. Auto-correction OCR (ocr_cleaner — optionnel) Exemple d'appel depuis Streamlit : from ocr.doctr_engine import ocr_pipeline result = ocr_pipeline(uploaded_file, language_hint="ar") words = result["words"] # liste LayoutXLM (corrigée) text = result["full_text"] # texte brut viz = result["visualization"] # PIL.Image annotée report = result["correction_report"] # DataFrame des corrections Args: file_input : UploadedFile Streamlit, chemin, bytes ou PIL.Image. language_hint : Code ISO optionnel ('fr', 'en', 'ar'). min_confidence : Seuil de confiance minimal. apply_cleaning : Active l'auto-correction post-OCR (défaut : True). Returns: dict : words, full_text, page_w, page_h, visualization, preprocessed, correction_report (DataFrame ou None). """ from preprocessing.image_preprocessor import preprocess_for_ocr # ── 1. Prétraitement ────────────────────────────────────────────────────── preprocessed_pil = preprocess_for_ocr( file_input, dpi=200, # 200 DPI suffisant (300 trop lourd sur CPU) denoise=True, deskew=True, binarize=True, return_pil=True, ) # ── 2. OCR ─────────────────────────────────────────────────────────────── ocr_result = run_ocr( preprocessed_pil, det_arch=_DEFAULT_DET, reco_arch=_DEFAULT_RECO, min_confidence=min_confidence, ) # ── 3. Visualisation ───────────────────────────────────────────────────── viz = draw_bounding_boxes( preprocessed_pil, ocr_result["words"], show_text=True, show_confidence=False, ) # ── 4. Auto-correction post-OCR ────────────────────────────────────────── words_final = ocr_result["words"] corr_report = None cleaning_error = None if apply_cleaning and words_final: try: import importlib _cleaner = importlib.import_module("ocr.ocr_cleaner") words_final = _cleaner.clean_ocr_output(words_final) corr_report = _cleaner.correction_report(words_final) logger.info( "Auto-correction : %d mot(s) corrigé(s) sur %d.", len(corr_report), len(words_final), ) except Exception as exc: cleaning_error = str(exc) logger.warning("Erreur auto-correction : %s", exc) # ── 5. Post-traitement métier (vocabulaire factures, mots fusionnés…) ──── post_text = None post_corr_report = None post_processing_error = None if words_final: try: import importlib _pp = importlib.import_module("ocr.post_processor") words_final = _pp.clean_invoice_ocr(words_final) lines = _pp.reconstruct_lines_table_aware(words_final) col_bounds = _pp.detect_column_boundaries(words_final) post_text = _pp.lines_to_text( lines, col_bounds if len(col_bounds) >= 1 else None, ) post_corr_report = _pp.correction_report_invoice(words_final) logger.info( "Post-traitement : %d correction(s), %d ligne(s), %d colonne(s).", len(post_corr_report), len(lines), len(col_bounds) + 1, ) except Exception as exc: post_processing_error = str(exc) logger.warning("Erreur post-traitement : %s", exc) # Nettoyage mémoire gc.collect() return { "words": words_final, "full_text": ocr_result["full_text"], "page_w": ocr_result["page_w"], "page_h": ocr_result["page_h"], "visualization": viz, "preprocessed": preprocessed_pil, "correction_report": corr_report, "cleaning_error": cleaning_error, "post_text": post_text, "post_corr_report": post_corr_report, "post_processing_error": post_processing_error, } # ───────────────────────────────────────────────────────────────────────────── # Utilitaires internes # ───────────────────────────────────────────────────────────────────────────── def _resize_for_inference(pil: Image.Image, max_width: int) -> Image.Image: """ Réduit l'image à max_width si elle est plus large. Conserve le ratio. Utilise LANCZOS (meilleur pour le texte). """ w, h = pil.size if w <= max_width: return pil new_h = int(h * max_width / w) return pil.resize((max_width, new_h), Image.LANCZOS) def _to_rgb_pil(image_input) -> Image.Image: """Convertit numpy (BGR/gris) ou PIL en PIL.Image RGB.""" if isinstance(image_input, Image.Image): return image_input.convert("RGB") import cv2 arr = np.asarray(image_input) if arr.ndim == 2: arr_rgb = cv2.cvtColor(arr, cv2.COLOR_GRAY2RGB) elif arr.shape[2] == 4: arr_rgb = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGB) else: arr_rgb = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB) return Image.fromarray(arr_rgb) def _is_arabic(text: str) -> bool: """Détecte si un mot contient majoritairement des caractères arabes (U+0600-U+06FF).""" if not text: return False arabic = sum(1 for ch in text if "\u0600" <= ch <= "\u06FF") return arabic / len(text) > 0.4 def _reconstruct_text(doctr_result) -> str: """Reconstruit le texte complet : mots → lignes → blocs.""" lines_out = [] for page in doctr_result.pages: for block in page.blocks: for line in block.lines: words = [w.value for w in line.words if w.value.strip()] if words: lines_out.append(" ".join(words)) return "\n".join(lines_out) class _NullContext: """Context manager vide si torch n'est pas disponible.""" def __enter__(self): return self def __exit__(self, *_): pass