4NK_IA_back/services/worker/pipelines/ocr.py

"""
Pipeline OCR pour l'extraction de texte
"""
import os
import logging
import subprocess
import tempfile
from utils.storage import store_artifact, cleanup_temp_file
from utils.text_normalize import correct_notarial_text

logger = logging.getLogger(__name__)

def run(doc_id: str, ctx: dict):
    """
    Étape OCR d'un document
    """
    logger.info(f"OCR du document {doc_id}")

    try:
        mime_type = ctx.get("mime_type", "application/pdf")

        if mime_type == "application/pdf":
            _ocr_pdf(doc_id, ctx)
        elif mime_type.startswith("image/"):
            _ocr_image(doc_id, ctx)
        else:
            raise ValueError(f"Type de fichier non supporté pour OCR: {mime_type}")

        # Stockage des métadonnées OCR
        ocr_meta = {
            "ocr_completed": True,
            "text_length": len(ctx.get("extracted_text", "")),
            "confidence": ctx.get("ocr_confidence", 0.0)
        }

        ctx["ocr_meta"] = ocr_meta

        logger.info(f"OCR terminé pour le document {doc_id}")

    except Exception as e:
        logger.error(f"Erreur lors de l'OCR du document {doc_id}: {e}")
        raise

def _ocr_pdf(doc_id: str, ctx: dict):
    """
    OCR spécifique aux PDF
    """
    try:
        temp_pdf = ctx.get("temp_pdf_path")
        if not temp_pdf:
            raise ValueError("Chemin du PDF temporaire non trouvé")

        pdf_meta = ctx.get("pdf_meta", {})

        # Si le PDF contient déjà du texte, l'extraire directement
        if pdf_meta.get("has_text", False):
            _extract_pdf_text(doc_id, ctx, temp_pdf)
        else:
            # OCR avec ocrmypdf
            _ocr_pdf_with_ocrmypdf(doc_id, ctx, temp_pdf)

    except Exception as e:
        logger.error(f"Erreur lors de l'OCR PDF pour {doc_id}: {e}")
        raise

def _extract_pdf_text(doc_id: str, ctx: dict, pdf_path: str):
    """
    Extraction de texte natif d'un PDF
    """
    try:
        import PyPDF2

        with open(pdf_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            text_parts = []

            for page_num, page in enumerate(pdf_reader.pages):
                page_text = page.extract_text()
                if page_text.strip():
                    text_parts.append(f"=== PAGE {page_num + 1} ===\n{page_text}")

            extracted_text = "\n\n".join(text_parts)

            # Correction lexicale
            corrected_text = correct_notarial_text(extracted_text)

            # Stockage du texte
            ctx["extracted_text"] = corrected_text
            ctx["ocr_confidence"] = 1.0  # Texte natif = confiance maximale

            # Stockage en artefact
            store_artifact(doc_id, "extracted_text.txt", corrected_text.encode('utf-8'), "text/plain")

            logger.info(f"Texte natif extrait du PDF {doc_id}: {len(corrected_text)} caractères")

    except Exception as e:
        logger.error(f"Erreur lors de l'extraction de texte natif pour {doc_id}: {e}")
        raise

def _ocr_pdf_with_ocrmypdf(doc_id: str, ctx: dict, pdf_path: str):
    """
    OCR d'un PDF avec ocrmypdf
    """
    try:
        # Création d'un fichier de sortie temporaire
        output_pdf = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
        output_txt = tempfile.NamedTemporaryFile(suffix=".txt", delete=False)
        output_pdf.close()
        output_txt.close()

        try:
            # Exécution d'ocrmypdf
            cmd = [
                "ocrmypdf",
                "--sidecar", output_txt.name,
                "--output-type", "pdf",
                "--language", "fra",
                "--optimize", "1",
                pdf_path,
                output_pdf.name
            ]

            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)

            if result.returncode != 0:
                raise RuntimeError(f"ocrmypdf a échoué: {result.stderr}")

            # Lecture du texte extrait
            with open(output_txt.name, 'r', encoding='utf-8') as f:
                extracted_text = f.read()

            # Correction lexicale
            corrected_text = correct_notarial_text(extracted_text)

            # Stockage du texte
            ctx["extracted_text"] = corrected_text
            ctx["ocr_confidence"] = 0.8  # Estimation pour OCR

            # Stockage des artefacts
            store_artifact(doc_id, "extracted_text.txt", corrected_text.encode('utf-8'), "text/plain")

            # Stockage du PDF OCRisé
            with open(output_pdf.name, 'rb') as f:
                ocr_pdf_content = f.read()
            store_artifact(doc_id, "ocr.pdf", ocr_pdf_content, "application/pdf")

            logger.info(f"OCR PDF terminé pour {doc_id}: {len(corrected_text)} caractères")

        finally:
            # Nettoyage des fichiers temporaires
            cleanup_temp_file(output_pdf.name)
            cleanup_temp_file(output_txt.name)

    except Exception as e:
        logger.error(f"Erreur lors de l'OCR PDF avec ocrmypdf pour {doc_id}: {e}")
        raise

def _ocr_image(doc_id: str, ctx: dict):
    """
    OCR d'une image avec Tesseract
    """
    try:
        temp_image = ctx.get("temp_image_path")
        if not temp_image:
            raise ValueError("Chemin de l'image temporaire non trouvé")

        import pytesseract
        from PIL import Image

        # Ouverture de l'image
        with Image.open(temp_image) as img:
            # Configuration Tesseract pour le français
            custom_config = r'--oem 3 --psm 6 -l fra'

            # Extraction du texte
            extracted_text = pytesseract.image_to_string(img, config=custom_config)

            # Récupération des données de confiance
            try:
                data = pytesseract.image_to_data(img, config=custom_config, output_type=pytesseract.Output.DICT)
                confidences = [int(conf) for conf in data['conf'] if int(conf) > 0]
                avg_confidence = sum(confidences) / len(confidences) / 100.0 if confidences else 0.0
            except:
                avg_confidence = 0.7  # Estimation par défaut

            # Correction lexicale
            corrected_text = correct_notarial_text(extracted_text)

            # Stockage du texte
            ctx["extracted_text"] = corrected_text
            ctx["ocr_confidence"] = avg_confidence

            # Stockage en artefact
            store_artifact(doc_id, "extracted_text.txt", corrected_text.encode('utf-8'), "text/plain")

            logger.info(f"OCR image terminé pour {doc_id}: {len(corrected_text)} caractères, confiance: {avg_confidence:.2f}")

    except Exception as e:
        logger.error(f"Erreur lors de l'OCR image pour {doc_id}: {e}")
        raise