Spaces:
Sleeping
Sleeping
| """ | |
| Pipeline de prétraitement d'images pour OCR multi-langue. | |
| Étapes : PDF/Image → Normalisation → Niveaux de gris → | |
| Débruitage → Deskewing → Binarisation adaptative | |
| Optimisé pour documents contenant du texte arabe cursif, | |
| français et anglais — photos de qualité variable acceptées. | |
| """ | |
| import io | |
| import logging | |
| from pathlib import Path | |
| from typing import Union, Optional | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| logger = logging.getLogger(__name__) | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| # Point d'entrée principal | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| def preprocess_for_ocr( | |
| file_input, | |
| dpi: int = 300, | |
| page_index: int = 0, | |
| denoise: bool = True, | |
| deskew: bool = True, | |
| binarize: bool = True, | |
| return_pil: bool = True, | |
| ) -> Union[Image.Image, np.ndarray]: | |
| """ | |
| Pipeline complet de prétraitement : charge le fichier puis applique | |
| niveaux de gris → débruitage → correction d'inclinaison → binarisation. | |
| Args: | |
| file_input : UploadedFile Streamlit, chemin (str/Path), bytes ou PIL.Image. | |
| dpi : Résolution pour la conversion PDF (300 recommandé pour l'OCR). | |
| page_index : Numéro de page à traiter (PDF multi-pages). | |
| denoise : Activer la réduction de bruit. | |
| deskew : Activer la correction automatique d'inclinaison. | |
| binarize : Activer la binarisation adaptative. | |
| return_pil : True → retourne PIL.Image ; False → numpy uint8. | |
| Returns: | |
| Image prétraitée prête pour docTR / PaddleOCR / EasyOCR. | |
| """ | |
| # ── Chargement ──────────────────────────────────────────────────────────── | |
| pages = _load_document(file_input, dpi=dpi) | |
| img = pages[min(page_index, len(pages) - 1)].copy() | |
| # ── 1. Normalisation de la résolution (upscale si trop petite) ─────────── | |
| img = _normalize_resolution(img, min_width=1200) | |
| # ── 2. Niveaux de gris ─────────────────────────────────────────────────── | |
| gray = _to_grayscale(img) | |
| # ── 3. Débruitage (paramètres conservateurs pour l'arabe cursif) ───────── | |
| if denoise: | |
| gray = _denoise(gray) | |
| # ── 4. Correction d'inclinaison ────────────────────────────────────────── | |
| if deskew: | |
| gray = _deskew(gray) | |
| # ── 5. Binarisation adaptative ─────────────────────────────────────────── | |
| result = _binarize(gray) if binarize else gray | |
| return _cv2_to_pil(result) if return_pil else result | |
| def get_preview_pair( | |
| file_input, | |
| dpi: int = 300, | |
| page_index: int = 0, | |
| ) -> tuple: | |
| """ | |
| Retourne (original_pil, preprocessed_pil) pour affichage côte à côte. | |
| Args: | |
| file_input : Source du fichier. | |
| dpi : Résolution DPI pour PDF. | |
| page_index : Page à utiliser. | |
| Returns: | |
| (PIL.Image original, PIL.Image prétraitée) | |
| """ | |
| pages = _load_document(file_input, dpi=dpi) | |
| raw = pages[min(page_index, len(pages) - 1)] | |
| original_pil = _cv2_to_pil(raw) | |
| preprocessed_pil = preprocess_for_ocr( | |
| file_input, dpi=dpi, page_index=page_index, return_pil=True | |
| ) | |
| return original_pil, preprocessed_pil | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| # Chargement multi-format | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| def _load_document(file_input, dpi: int = 300) -> list: | |
| """ | |
| Charge un document et retourne une liste de pages (numpy BGR). | |
| Formats supportés : JPEG, PNG, TIFF, BMP, PDF. | |
| """ | |
| # PIL.Image direct | |
| if isinstance(file_input, Image.Image): | |
| return [_pil_to_cv2(file_input)] | |
| # numpy array direct | |
| if isinstance(file_input, np.ndarray): | |
| return [file_input] | |
| # Lire les bytes selon la source | |
| if hasattr(file_input, "read"): | |
| # UploadedFile Streamlit ou file-like object | |
| raw_bytes = file_input.read() | |
| filename = getattr(file_input, "name", "file.bin") | |
| try: | |
| file_input.seek(0) # Rembobiner pour usage ultérieur dans l'app | |
| except Exception: | |
| pass | |
| elif isinstance(file_input, (str, Path)): | |
| filename = str(file_input) | |
| with open(file_input, "rb") as f: | |
| raw_bytes = f.read() | |
| elif isinstance(file_input, bytes): | |
| raw_bytes = file_input | |
| filename = "file.bin" | |
| else: | |
| raise TypeError(f"Type d'entrée non supporté : {type(file_input)}") | |
| ext = Path(filename).suffix.lower() | |
| # ── PDF → images haute résolution ───────────────────────────────────────── | |
| if ext == ".pdf": | |
| try: | |
| from pdf2image import convert_from_bytes | |
| pil_pages = convert_from_bytes(raw_bytes, dpi=dpi) | |
| return [_pil_to_cv2(p) for p in pil_pages] | |
| except ImportError: | |
| raise ImportError( | |
| "pdf2image est requis pour les PDF.\n" | |
| "Installez-le avec : pip install pdf2image\n" | |
| "Sous Windows, téléchargez aussi Poppler : " | |
| "https://github.com/oschwartz10612/poppler-windows/releases" | |
| ) | |
| # ── Image standard (JPEG, PNG, TIFF, BMP…) ──────────────────────────────── | |
| arr = np.frombuffer(raw_bytes, dtype=np.uint8) | |
| img = cv2.imdecode(arr, cv2.IMREAD_COLOR) | |
| if img is None: | |
| raise ValueError( | |
| f"Impossible de décoder l'image '{filename}'. " | |
| "Vérifiez que le fichier n'est pas corrompu." | |
| ) | |
| return [img] | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| # Étapes du pipeline | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| def _normalize_resolution( | |
| img: np.ndarray, | |
| min_width: int = 1400, | |
| max_width: int = 2400, | |
| ) -> np.ndarray: | |
| """ | |
| Maintient la largeur de l'image dans [min_width, max_width]. | |
| - Upscale si trop petite (< min_width) : améliore la précision OCR. | |
| · LANCZOS4 si scale > 1.5 (meilleure conservation des arêtes de chars) | |
| · CUBIC sinon (plus rapide, suffisant pour facteur faible) | |
| - Downscale si trop grande (> max_width) : réduit la RAM et le temps | |
| de traitement sur CPU — INTER_AREA est optimal pour réduire. | |
| Seuils relevés à 1400/2400 px (vs 1000/1800) pour réduire le CER | |
| sur les factures avec petits caractères (numéros de ligne, montants). | |
| Consommation mémoire : ≈ 280 Mo max sur CPU 8 Go. | |
| """ | |
| h, w = img.shape[:2] | |
| if w < min_width: | |
| scale = min_width / w | |
| interp = cv2.INTER_LANCZOS4 if scale > 1.5 else cv2.INTER_CUBIC | |
| img = cv2.resize( | |
| img, | |
| (int(w * scale), int(h * scale)), | |
| interpolation=interp, | |
| ) | |
| elif w > max_width: | |
| scale = max_width / w | |
| img = cv2.resize( | |
| img, | |
| (int(w * scale), int(h * scale)), | |
| interpolation=cv2.INTER_AREA, | |
| ) | |
| return img | |
| def _to_grayscale(img: np.ndarray) -> np.ndarray: | |
| """Convertit BGR → niveaux de gris (no-op si déjà en gris).""" | |
| if len(img.shape) == 2: | |
| return img | |
| return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| def _denoise(gray: np.ndarray) -> np.ndarray: | |
| """ | |
| Réduction du bruit par filtre Non-Local Means. | |
| Paramètres conservateurs (h=7) pour préserver : | |
| - Les connexions entre lettres arabes (cursif) | |
| - Les diacritiques (harakat, shadda, tanwin…) | |
| - Les petits détails des caractères français accentués | |
| Un h trop élevé (>15) estompe ces détails et dégrade l'OCR arabe. | |
| """ | |
| return cv2.fastNlMeansDenoising( | |
| gray, | |
| h=7, # Force du filtre — conservateur pour l'arabe | |
| templateWindowSize=7, | |
| searchWindowSize=21, | |
| ) | |
| def _deskew(gray: np.ndarray) -> np.ndarray: | |
| """ | |
| Correction automatique de l'inclinaison du document. | |
| Algorithme : | |
| 1. Binarisation Otsu temporaire (rapide) | |
| 2. Détection des lignes de texte via Hough probabiliste | |
| 3. Angle médian des lignes quasi-horizontales | |
| 4. Rotation de correction si l'angle est > 0.3° (seuil de tolérance) | |
| Fonctionne pour l'arabe car les lignes de base sont horizontales | |
| quelle que soit la direction d'écriture (RTL). | |
| """ | |
| angle = _estimate_skew_angle(gray) | |
| if abs(angle) < 0.3: | |
| return gray # Inclinaison négligeable → pas de traitement | |
| logger.debug("Deskew : correction de %.2f°", angle) | |
| return _rotate_image(gray, angle) | |
| def _estimate_skew_angle(gray: np.ndarray) -> float: | |
| """ | |
| Estime l'angle d'inclinaison dominant (en degrés). | |
| Retourne 0.0 si aucune ligne exploitable n'est détectée. | |
| """ | |
| # Binarisation grossière pour la détection de contours | |
| _, binary = cv2.threshold( | |
| gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU | |
| ) | |
| edges = cv2.Canny(binary, 50, 150, apertureSize=3) | |
| # Longueur minimale de ligne = 25 % de la largeur (ignore le bruit) | |
| min_len = gray.shape[1] // 4 | |
| lines = cv2.HoughLinesP( | |
| edges, | |
| rho=1, | |
| theta=np.pi / 180, | |
| threshold=80, | |
| minLineLength=min_len, | |
| maxLineGap=20, | |
| ) | |
| if lines is None or len(lines) == 0: | |
| return 0.0 | |
| angles = [] | |
| for line in lines: | |
| x1, y1, x2, y2 = line[0] | |
| dx = x2 - x1 | |
| if dx == 0: | |
| continue | |
| angle = np.degrees(np.arctan2(y2 - y1, dx)) | |
| # Garder uniquement les lignes proches de l'horizontale (±15°) | |
| if -15.0 < angle < 15.0: | |
| angles.append(angle) | |
| if not angles: | |
| return 0.0 | |
| # Médiane pour être robuste face aux lignes aberrantes | |
| return float(np.median(angles)) | |
| def _rotate_image(gray: np.ndarray, angle: float) -> np.ndarray: | |
| """ | |
| Rotation autour du centre de l'image, fond blanc. | |
| Les nouvelles dimensions sont recalculées pour éviter le rognage. | |
| """ | |
| h, w = gray.shape[:2] | |
| cx, cy = w // 2, h // 2 | |
| M = cv2.getRotationMatrix2D((cx, cy), angle, scale=1.0) | |
| cos_a = abs(M[0, 0]) | |
| sin_a = abs(M[0, 1]) | |
| new_w = int(h * sin_a + w * cos_a) | |
| new_h = int(h * cos_a + w * sin_a) | |
| # Ajuster la translation pour centrer l'image dans les nouvelles dimensions | |
| M[0, 2] += (new_w - w) / 2 | |
| M[1, 2] += (new_h - h) / 2 | |
| return cv2.warpAffine( | |
| gray, M, (new_w, new_h), | |
| flags=cv2.INTER_CUBIC, | |
| borderMode=cv2.BORDER_CONSTANT, | |
| borderValue=255, # Fond blanc | |
| ) | |
| def _binarize(gray: np.ndarray) -> np.ndarray: | |
| """ | |
| Binarisation adaptative gaussienne. | |
| Avantages sur Otsu global : | |
| - Gère les ombres portées et l'éclairage non uniforme (photos de documents) | |
| - Préserve les zones de texte foncé sur fond clair ET clair sur fond foncé | |
| - blockSize=31 : fenêtre suffisamment grande pour les ombres diffuses | |
| - C=11 : soustraction de constante locale pour affiner le seuil | |
| Pour l'arabe : | |
| - La binarisation adaptative respecte mieux les diacritiques (points, | |
| shadda, tanwin) que la binarisation globale qui les efface parfois. | |
| - Pré-lissage léger (GaussianBlur 3×3) pour éviter l'amplification | |
| du bruit dans les zones uniformes avant seuillage. | |
| """ | |
| # Lissage préalable pour réduire le bruit haute fréquence | |
| blurred = cv2.GaussianBlur(gray, (3, 3), 0) | |
| binary = cv2.adaptiveThreshold( | |
| blurred, | |
| maxValue=255, | |
| adaptiveMethod=cv2.ADAPTIVE_THRESH_GAUSSIAN_C, | |
| thresholdType=cv2.THRESH_BINARY, | |
| blockSize=31, | |
| C=11, | |
| ) | |
| return binary | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| # Utilitaires de conversion PIL ↔ OpenCV | |
| # ───────────────────────────────────────────────────────────────────────────── | |
| def _pil_to_cv2(pil_img: Image.Image) -> np.ndarray: | |
| """PIL.Image (RGB) → numpy BGR pour OpenCV.""" | |
| rgb = np.array(pil_img.convert("RGB")) | |
| return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) | |
| def _cv2_to_pil(img: np.ndarray) -> Image.Image: | |
| """numpy BGR ou niveaux de gris → PIL.Image.""" | |
| if len(img.shape) == 2: | |
| return Image.fromarray(img, mode="L") | |
| rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| return Image.fromarray(rgb) | |