mirror of
https://github.com/Ladebeze66/ragflow_preprocess.git
synced 2026-02-04 05:50:26 +01:00
271 lines
11 KiB
Python
271 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
Module OCR pour la reconnaissance de texte dans les images
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytesseract
|
|
from typing import Union, Dict, Tuple, Optional
|
|
from PIL import Image
|
|
import io
|
|
import copy
|
|
import os
|
|
|
|
# Configuration du chemin de Tesseract OCR
|
|
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe' # Pour Windows
|
|
# Pour Linux et macOS, Tesseract doit être installé et disponible dans le PATH
|
|
|
|
# Vérification si le chemin de Tesseract existe, sinon essayer d'autres chemins courants
|
|
if not os.path.exists(pytesseract.pytesseract.tesseract_cmd):
|
|
alternative_paths = [
|
|
r'C:\Program Files (x86)\Tesseract-OCR\tesseract.exe',
|
|
r'C:\Tesseract-OCR\tesseract.exe'
|
|
]
|
|
for path in alternative_paths:
|
|
if os.path.exists(path):
|
|
pytesseract.pytesseract.tesseract_cmd = path
|
|
break
|
|
|
|
class OCRProcessor:
|
|
"""
|
|
Classe pour traiter les images et extraire le texte
|
|
"""
|
|
|
|
def __init__(self, lang: str = "fra+eng"):
|
|
"""
|
|
Initialise le processeur OCR
|
|
|
|
Args:
|
|
lang (str): Langues pour la reconnaissance (fra pour français, eng pour anglais)
|
|
"""
|
|
self.lang = lang
|
|
|
|
def preprocess_image(self, image: Union[bytes, np.ndarray, str]) -> np.ndarray:
|
|
"""
|
|
Prétraite l'image pour améliorer la reconnaissance OCR
|
|
|
|
Args:
|
|
image: Image sous forme de bytes, numpy array ou chemin de fichier
|
|
|
|
Returns:
|
|
np.ndarray: Image prétraitée
|
|
"""
|
|
# Convertir l'image en numpy array si nécessaire
|
|
if isinstance(image, bytes):
|
|
img = np.array(Image.open(io.BytesIO(image)))
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
elif isinstance(image, str):
|
|
img = cv2.imread(image)
|
|
else:
|
|
# Pour un np.ndarray, on peut utiliser copy() directement
|
|
# Pour d'autres types, on s'assure de créer une copie sécurisée
|
|
img = np.array(image) if not isinstance(image, np.ndarray) else image.copy()
|
|
|
|
# Convertir en niveaux de gris
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
|
|
# Appliquer une réduction de bruit
|
|
denoised = cv2.fastNlMeansDenoising(gray, None, 10, 7, 21)
|
|
|
|
# Seuillage adaptatif
|
|
binary = cv2.adaptiveThreshold(
|
|
denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
|
cv2.THRESH_BINARY, 11, 2
|
|
)
|
|
|
|
return binary
|
|
|
|
def extract_text(self, image: Union[bytes, np.ndarray, str],
|
|
preprocess: bool = True) -> str:
|
|
"""
|
|
Extrait le texte d'une image
|
|
|
|
Args:
|
|
image: Image sous forme de bytes, numpy array ou chemin de fichier
|
|
preprocess (bool): Si True, prétraite l'image avant la reconnaissance
|
|
|
|
Returns:
|
|
str: Texte extrait de l'image
|
|
"""
|
|
try:
|
|
# Prétraitement si demandé
|
|
if preprocess:
|
|
processed_img = self.preprocess_image(image)
|
|
else:
|
|
if isinstance(image, bytes):
|
|
processed_img = np.array(Image.open(io.BytesIO(image)))
|
|
elif isinstance(image, str):
|
|
processed_img = cv2.imread(image)
|
|
else:
|
|
processed_img = np.array(image) if not isinstance(image, np.ndarray) else image
|
|
|
|
# Lancer la reconnaissance OCR
|
|
config = r'--oem 3 --psm 6' # Page entière en mode 6
|
|
text = pytesseract.image_to_string(processed_img, lang=self.lang, config=config)
|
|
|
|
return text.strip()
|
|
|
|
except Exception as e:
|
|
print(f"Erreur lors de l'extraction de texte: {str(e)}")
|
|
return ""
|
|
|
|
def extract_text_with_confidence(self, image: Union[bytes, np.ndarray, str],
|
|
preprocess: bool = True) -> Dict:
|
|
"""
|
|
Extrait le texte avec les données de confiance
|
|
|
|
Args:
|
|
image: Image sous forme de bytes, numpy array ou chemin de fichier
|
|
preprocess (bool): Si True, prétraite l'image avant la reconnaissance
|
|
|
|
Returns:
|
|
Dict: Dictionnaire contenant le texte et les données de confiance
|
|
"""
|
|
try:
|
|
# Prétraitement si demandé
|
|
if preprocess:
|
|
processed_img = self.preprocess_image(image)
|
|
else:
|
|
if isinstance(image, bytes):
|
|
processed_img = np.array(Image.open(io.BytesIO(image)))
|
|
elif isinstance(image, str):
|
|
processed_img = cv2.imread(image)
|
|
else:
|
|
processed_img = np.array(image) if not isinstance(image, np.ndarray) else image
|
|
|
|
# Lancer la reconnaissance OCR avec données détaillées
|
|
config = r'--oem 3 --psm 6'
|
|
data = pytesseract.image_to_data(processed_img, lang=self.lang, config=config, output_type=pytesseract.Output.DICT)
|
|
|
|
# Calculer la confiance moyenne
|
|
confidences = [conf for conf in data['conf'] if conf != -1]
|
|
avg_confidence = sum(confidences) / len(confidences) if confidences else 0
|
|
|
|
# Reconstruire le texte à partir des mots reconnus
|
|
text = ' '.join([word for word in data['text'] if word.strip()])
|
|
|
|
return {
|
|
'text': text,
|
|
'confidence': avg_confidence,
|
|
'words': data['text'],
|
|
'word_confidences': data['conf']
|
|
}
|
|
|
|
except Exception as e:
|
|
print(f"Erreur lors de l'extraction de texte avec confiance: {str(e)}")
|
|
return {'text': "", 'confidence': 0, 'words': [], 'word_confidences': []}
|
|
|
|
def detect_tables(self, image: Union[bytes, np.ndarray, str]) -> Optional[np.ndarray]:
|
|
"""
|
|
Détecte les tableaux dans une image
|
|
|
|
Args:
|
|
image: Image sous forme de bytes, numpy array ou chemin de fichier
|
|
|
|
Returns:
|
|
Optional[np.ndarray]: Image avec les tableaux identifiés ou None en cas d'erreur
|
|
"""
|
|
try:
|
|
# Convertir l'image en numpy array si nécessaire
|
|
if isinstance(image, bytes):
|
|
img = np.array(Image.open(io.BytesIO(image)))
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
elif isinstance(image, str):
|
|
img = cv2.imread(image)
|
|
else:
|
|
img = np.array(image) if not isinstance(image, np.ndarray) else image.copy()
|
|
|
|
# Convertir en niveaux de gris
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
|
|
# Seuillage
|
|
thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
|
cv2.THRESH_BINARY_INV, 11, 2)
|
|
|
|
# Dilatation pour renforcer les lignes
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
|
dilate = cv2.dilate(thresh, kernel, iterations=3)
|
|
|
|
# Trouver les contours
|
|
contours, _ = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
# Dessiner les contours des tableaux potentiels sur une copie de l'image
|
|
result = img.copy()
|
|
|
|
for contour in contours:
|
|
area = cv2.contourArea(contour)
|
|
if area > 1000: # Filtrer les petits contours
|
|
x, y, w, h = cv2.boundingRect(contour)
|
|
# Un tableau a généralement un ratio largeur/hauteur spécifique
|
|
if 0.5 < w/h < 2: # Ajuster selon les besoins
|
|
cv2.rectangle(result, (x, y), (x+w, y+h), (0, 255, 0), 2)
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
print(f"Erreur lors de la détection de tableaux: {str(e)}")
|
|
return None
|
|
|
|
def extract_table_data(self, image: Union[bytes, np.ndarray, str]) -> Optional[Dict]:
|
|
"""
|
|
Extrait les données d'un tableau détecté dans l'image
|
|
|
|
Args:
|
|
image: Image sous forme de bytes, numpy array ou chemin de fichier
|
|
|
|
Returns:
|
|
Optional[Dict]: Dictionnaire contenant les données du tableau ou None en cas d'erreur
|
|
"""
|
|
# Cette fonction est une simplification. Pour une reconnaissance complète de tableaux,
|
|
# des librairies spécialisées comme camelot-py ou tabula-py seraient plus appropriées.
|
|
|
|
try:
|
|
# Pour cette démonstration, on utilise pytesseract avec un mode de segmentation adapté aux tableaux
|
|
if isinstance(image, bytes):
|
|
img = np.array(Image.open(io.BytesIO(image)))
|
|
elif isinstance(image, str):
|
|
img = cv2.imread(image)
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
else:
|
|
# Conversion sécurisée pour tout type d'entrée
|
|
img_array = np.array(image)
|
|
img = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
|
|
|
|
# Configuration pour tableaux
|
|
config = r'--oem 3 --psm 6 -c preserve_interword_spaces=1'
|
|
data = pytesseract.image_to_data(img, lang=self.lang, config=config, output_type=pytesseract.Output.DICT)
|
|
|
|
# Organiser les données en lignes et colonnes (simplification)
|
|
words = data['text']
|
|
left = data['left']
|
|
top = data['top']
|
|
|
|
# Identifier les lignes uniques basées sur la position verticale
|
|
unique_tops = sorted(list(set([t for t, w in zip(top, words) if w.strip()])))
|
|
rows = []
|
|
|
|
for t in unique_tops:
|
|
# Trouver tous les mots sur cette ligne (avec une tolérance)
|
|
tolerance = 10
|
|
row_words = [(l, w) for l, w, wt in zip(left, words, top)
|
|
if wt >= t-tolerance and wt <= t+tolerance and w.strip()]
|
|
|
|
# Trier les mots par position horizontale
|
|
row_words.sort(key=lambda x: x[0])
|
|
|
|
# Ajouter les mots de cette ligne
|
|
rows.append([w for _, w in row_words])
|
|
|
|
# Construire un dictionnaire de résultats
|
|
return {
|
|
'rows': rows,
|
|
'raw_text': ' '.join([word for word in words if word.strip()]),
|
|
'num_rows': len(rows)
|
|
}
|
|
|
|
except Exception as e:
|
|
print(f"Erreur lors de l'extraction des données du tableau: {str(e)}")
|
|
return None |