scan-ocr-cheques/worker/app/tasks.py

"""
Tâches de traitement pour le worker Cheque Scanner
"""

import os
import sys
import time
import json
import logging
import traceback
import redis
from datetime import datetime
from rq import get_current_job

# Ajouter le module d'extraction au path
sys.path.append('/app/shared')

# Importer les fonctions d'extraction
from extraction import (
    extraire_infos_cheque,
    get_tessdata_path
)

# Configuration du logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("cheque_scanner_tasks")

# Variables d'environnement
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
DEFAULT_OCR_LANGUAGE = os.getenv("DEFAULT_OCR_LANGUAGE", "eng")
ALTERNATIVE_OCR_LANGUAGE = os.getenv("ALTERNATIVE_OCR_LANGUAGE", "fra")
TESSERACT_DATA_PATH = os.getenv("TESSERACT_DATA_PATH", "/usr/share/tesseract-ocr/4.00/tessdata")


def update_job_status(job_id, status, message=None, progress=None, result=None, texte_brut=None, erreur=None, methode=None):
    """
    Met à jour le statut d'une tâche dans Redis
    """
    try:
        # Connexion à Redis
        redis_conn = redis.Redis.from_url(REDIS_URL)
        
        # Préparer les données à mettre à jour
        update_data = {
            "status": status,
            "updated_at": datetime.now().isoformat()
        }
        
        if message:
            update_data["message"] = message
        
        if progress:
            update_data["progress"] = str(progress)
        
        if result:
            update_data["result"] = str(result)
        
        if texte_brut:
            update_data["texte_brut"] = texte_brut
        
        if erreur:
            update_data["erreur"] = erreur
        
        if methode:
            update_data["methode"] = methode
        
        if status == "completed":
            update_data["completed_at"] = datetime.now().isoformat()
        
        # Mettre à jour les données dans Redis
        redis_conn.hset(f"job:{job_id}", mapping=update_data)
        
        logger.info(f"Statut de la tâche {job_id} mis à jour: {status}")
        
        return True
    
    except Exception as e:
        logger.error(f"Erreur lors de la mise à jour du statut de la tâche {job_id}: {str(e)}")
        return False


def process_cheque_image(file_path, job_id):
    """
    Traite une image de chèque pour en extraire les informations
    
    Args:
        file_path (str): Chemin vers l'image à traiter
        job_id (str): Identifiant de la tâche
    
    Returns:
        dict: Résultat de l'extraction
    """
    job = get_current_job()
    logger.info(f"Début du traitement de l'image: {file_path} (Tâche: {job_id})")
    
    # Mettre à jour le statut
    update_job_status(
        job_id=job_id,
        status="processing",
        message="Traitement en cours",
        progress=10
    )
    
    try:
        # Vérifier que le fichier existe
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"L'image {file_path} n'existe pas")
        
        # Récupérer le chemin vers tessdata
        tessdata_path = TESSERACT_DATA_PATH
        if not os.path.exists(tessdata_path):
            # Essayer de trouver automatiquement
            tessdata_path = get_tessdata_path()
        
        # Mise à jour intermédiaire
        update_job_status(
            job_id=job_id,
            status="processing",
            message="Extraction des informations en cours",
            progress=30
        )
        
        # Première tentative avec la langue par défaut
        try:
            logger.info(f"Tentative d'extraction avec la langue: {DEFAULT_OCR_LANGUAGE}")
            update_job_status(
                job_id=job_id,
                status="processing",
                message=f"Extraction avec langue {DEFAULT_OCR_LANGUAGE}",
                progress=50
            )
            
            infos, texte = extraire_infos_cheque(
                chemin_image=file_path,
                methode="ocr",
                language=DEFAULT_OCR_LANGUAGE,
                tessdata=tessdata_path
            )
            
            methode = "ocr"
            
        except Exception as e:
            logger.warning(f"Échec de la première tentative: {str(e)}")
            
            # Deuxième tentative avec la langue alternative
            try:
                logger.info(f"Tentative d'extraction avec la langue: {ALTERNATIVE_OCR_LANGUAGE}")
                update_job_status(
                    job_id=job_id,
                    status="processing",
                    message=f"Extraction avec langue {ALTERNATIVE_OCR_LANGUAGE}",
                    progress=60
                )
                
                infos, texte = extraire_infos_cheque(
                    chemin_image=file_path,
                    methode="ocr",
                    language=ALTERNATIVE_OCR_LANGUAGE,
                    tessdata=tessdata_path
                )
                
                methode = "ocr"
                
            except Exception as e2:
                logger.warning(f"Échec de la deuxième tentative: {str(e2)}")
                
                # Troisième tentative avec la méthode CV
                logger.info("Tentative d'extraction avec la méthode CV (sans OCR)")
                update_job_status(
                    job_id=job_id,
                    status="processing",
                    message="Extraction par détection de zones (sans OCR)",
                    progress=70
                )
                
                infos, texte = extraire_infos_cheque(
                    chemin_image=file_path,
                    methode="cv"
                )
                
                methode = "cv"
        
        # Mise à jour finale
        update_job_status(
            job_id=job_id,
            status="processing",
            message="Finalisation des résultats",
            progress=90
        )
        
        # Sauvegarder les résultats dans Redis
        update_job_status(
            job_id=job_id,
            status="completed",
            message="Extraction terminée avec succès",
            progress=100,
            result=infos,
            texte_brut=texte,
            methode=methode
        )
        
        logger.info(f"Traitement terminé pour la tâche {job_id}")
        
        return {
            "job_id": job_id,
            "status": "completed",
            "result": infos,
            "methode": methode
        }
    
    except Exception as e:
        # Capturer l'erreur
        error_trace = traceback.format_exc()
        logger.error(f"Erreur lors du traitement de l'image: {str(e)}\n{error_trace}")
        
        # Mettre à jour le statut avec l'erreur
        update_job_status(
            job_id=job_id,
            status="failed",
            message="Échec du traitement",
            erreur=str(e)
        )
        
        # Retourner l'erreur
        return {
            "job_id": job_id,
            "status": "failed",
            "error": str(e)
        }