llm_ticket3/utils/ocr_brut/ocr_utils.py
2025-05-07 17:12:50 +02:00

104 lines
3.8 KiB
Python

import pytesseract
import cv2
import numpy as np
import os
from pathlib import Path
from PIL import Image
from langdetect import detect
import re
from utils.ocr_brut.ocr_preprocessor import preprocess_image_with_profile
# ⬇️ PARAMÈTRES CENTRAUX D'ACTIVATION ⬇️
USE_PREPROCESSING = True # Active le prétraitement de l'image via ocr_preprocessor.py
USE_TEXT_CORRECTION = True # Corrige les mots tronqués après OCR
SAVE_DEBUG_OUTPUT = True # Sauvegarde image + texte dans debug_ocr/
AUTO_DETECT_LANGUAGE = True # Détecte automatiquement la langue
# Paramètres Tesseract optimaux selon tests
OCR_DEFAULT_PSM = 11 # Page segmentation mode 11 (sparse text)
OCR_DEFAULT_OEM = 3 # OCR Engine mode 3 (default, based on what is available)
OCR_DEFAULT_PROFILE = "document" # Profil de prétraitement optimal selon tests
# Complétion de mots tronqués (rudimentaire mais utile)
def completer_mots_tronques(texte):
lignes = texte.splitlines()
lignes_corrigees = []
for ligne in lignes:
if ligne.strip().endswith("-"):
ligne = ligne.strip()[:-1]
lignes_corrigees.append(ligne)
return "\n".join(lignes_corrigees)
# Détection de langue automatique (si activée)
def detect_language_tesseract(image_cv):
try:
text_sample = pytesseract.image_to_string(image_cv, config="--psm 6")
lang = detect(text_sample)
return {
"fr": "fra",
"en": "eng"
}.get(lang, "fra+eng")
except:
return "fra+eng"
# OCR principal
def extraire_texte(image_path, lang="auto"):
# Vérification de l'existence de l'image
if not os.path.exists(image_path):
print(f"[OCR] Image non trouvée: {image_path}")
return "", None
# Prétraitement de l'image avec le profil optimal
if USE_PREPROCESSING:
img_optimized = preprocess_image_with_profile(image_path, profile_name=OCR_DEFAULT_PROFILE)
else:
img_optimized = Image.open(image_path)
# Détection de langue
ocr_lang = lang
if lang == "auto" and AUTO_DETECT_LANGUAGE:
ocr_lang = detect_language_tesseract(img_optimized)
if ocr_lang == "auto":
ocr_lang = "fra+eng"
# OCR avec paramètres optimaux
config = f"--psm {OCR_DEFAULT_PSM} --oem {OCR_DEFAULT_OEM} -l {ocr_lang}"
texte = pytesseract.image_to_string(img_optimized, config=config)
# Correction des mots tronqués
if USE_TEXT_CORRECTION:
texte_corrige = completer_mots_tronques(texte)
if len(texte_corrige) >= len(texte) * 0.9:
texte = texte_corrige
# Sauvegarde debug (si activée)
if SAVE_DEBUG_OUTPUT and texte:
try:
debug_dir = "results/ocr_brut"
os.makedirs(debug_dir, exist_ok=True)
image_name = Path(image_path).stem
# Conversion si image PIL
if isinstance(img_optimized, Image.Image):
img_optimized = np.array(img_optimized)
if img_optimized.ndim == 3 and img_optimized.shape[2] == 3:
img_optimized = cv2.cvtColor(img_optimized, cv2.COLOR_RGB2BGR)
elif img_optimized.ndim == 3 and img_optimized.shape[2] == 4:
img_optimized = cv2.cvtColor(img_optimized, cv2.COLOR_RGBA2BGR)
if isinstance(img_optimized, np.ndarray):
cv2.imwrite(f"{debug_dir}/optimized_{image_name}.png", img_optimized)
with open(f"{debug_dir}/ocr_{image_name}.png.txt", "w", encoding="utf-8") as f:
f.write(texte)
except Exception as e:
print(f"[OCR DEBUG] Erreur de sauvegarde debug: {e}")
return texte, img_optimized
# Raccourci rapide pour juste récupérer le texte en français
def extraire_texte_fr(image_path):
texte, _ = extraire_texte(image_path, lang="fra")
return texte