llm_ticket3/test_ocr_to_image_sorter.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import json
import time
from PIL import Image
import pytesseract
from utils.ocr_utils import extraire_texte_fr, pretraiter_image
from utils.translate_utils import fr_to_en, en_to_fr
from llm_classes.llama_vision import LlamaVision
from agents.llama_vision.agent_image_sorter import AgentImageSorter

def parse_args():
    parser = argparse.ArgumentParser(description="Test complet OCR + Traduction + Tri d'images")
    parser.add_argument("image_path", help="Chemin vers l'image à analyser")
    parser.add_argument("--verbose", "-v", action="store_true", help="Mode verbeux")
    parser.add_argument("--only-ocr", "-o", action="store_true", help="Tester uniquement l'OCR")
    parser.add_argument("--skip-llm", "-s", action="store_true", help="Sauter l'appel au LLM (pour tester rapidement OCR)")
    parser.add_argument("--test-dir", "-d", help="Tester toutes les images d'un répertoire")
    return parser.parse_args()

def test_image(image_path, verbose=False, only_ocr=False, skip_llm=False):
    """
    Teste la chaîne de traitement complète sur une image
    """
    start_time = time.time()
    
    print(f"\n=== TEST DE L'IMAGE: {os.path.basename(image_path)} ===")
    
    if not os.path.exists(image_path):
        print(f"Erreur: Le fichier {image_path} n'existe pas")
        return None
    
    # Afficher les infos sur l'image
    try:
        with Image.open(image_path) as img:
            print(f"Format: {img.format}")
            print(f"Mode: {img.mode}")
            print(f"Taille: {img.size}")
    except Exception as e:
        print(f"Erreur lors de l'analyse de l'image: {e}")
        return None
    
    # ÉTAPE 1: OCR
    print("\n=== ÉTAPE 1: OCR ===")
    ocr_time_start = time.time()
    
    # Test avec l'image originale
    print("OCR sur image originale...")
    ocr_fr_original = ""
    try:
        with Image.open(image_path) as img:
            ocr_fr_original = pytesseract.image_to_string(img, lang="fra").strip()
    except Exception as e:
        print(f"Erreur OCR image originale: {e}")
    
    if ocr_fr_original:
        print(f"Texte détecté (original): {ocr_fr_original}")
    else:
        print("Aucun texte détecté sur l'image originale")
    
    # Test avec l'image prétraitée
    print("\nOCR avec image prétraitée...")
    try:
        # Utiliser notre fonction de prétraitement
        ocr_fr = extraire_texte_fr(image_path)
        if ocr_fr:
            print(f"Texte détecté (prétraité): {ocr_fr}")
        else:
            print("Aucun texte détecté après prétraitement")
    except Exception as e:
        print(f"Erreur lors de l'OCR prétraité: {e}")
        ocr_fr = ""
    
    ocr_time = time.time() - ocr_time_start
    print(f"Temps OCR: {ocr_time:.2f} secondes")
    
    # Si on s'arrête à l'OCR
    if only_ocr:
        total_time = time.time() - start_time
        print(f"\nTemps total: {total_time:.2f} secondes")
        return {
            "image": os.path.basename(image_path),
            "ocr_fr_original": ocr_fr_original,
            "ocr_fr": ocr_fr,
            "processing_time": {
                "ocr": ocr_time,
                "total": total_time
            }
        }
    
    # ÉTAPE 2: TRADUCTION
    print("\n=== ÉTAPE 2: TRADUCTION ===")
    translation_time_start = time.time()
    
    ocr_en = ""
    if ocr_fr:
        try:
            ocr_en = fr_to_en(ocr_fr)
            print(f"Traduction EN: {ocr_en}")
            
            # Traduction de vérification
            ocr_en_back_fr = en_to_fr(ocr_en)
            print(f"Traduction retour FR: {ocr_en_back_fr}")
        except Exception as e:
            print(f"Erreur lors de la traduction: {e}")
    else:
        print("Aucun texte à traduire")
    
    translation_time = time.time() - translation_time_start
    print(f"Temps traduction: {translation_time:.2f} secondes")
    
    # Si on saute l'appel au LLM
    if skip_llm:
        total_time = time.time() - start_time
        print(f"\nTemps total: {total_time:.2f} secondes")
        return {
            "image": os.path.basename(image_path),
            "ocr_fr": ocr_fr,
            "ocr_en": ocr_en,
            "processing_time": {
                "ocr": ocr_time,
                "translation": translation_time,
                "total": total_time
            }
        }
    
    # ÉTAPE 3: ANALYSE LLM
    print("\n=== ÉTAPE 3: ANALYSE LLM ===")
    llm_time_start = time.time()
    
    llm_result = None
    try:
        # Initialiser le modèle LlamaVision
        llm = LlamaVision()
        
        # Initialiser l'agent
        agent = AgentImageSorter(llm)
        
        # Exécuter l'agent
        llm_result = agent.executer(image_path)
        
        if llm_result:
            print(f"Image pertinente: {llm_result['is_relevant']}")
            print(f"Raison: {llm_result['reason']}")
        else:
            print("Erreur: Aucun résultat retourné par l'agent")
    except Exception as e:
        print(f"Erreur lors de l'analyse LLM: {e}")
    
    llm_time = time.time() - llm_time_start
    print(f"Temps LLM: {llm_time:.2f} secondes")
    
    # Temps total
    total_time = time.time() - start_time
    print(f"\nTemps total: {total_time:.2f} secondes")
    
    # Résultat complet
    result = {
        "image": os.path.basename(image_path),
        "ocr_fr": ocr_fr,
        "ocr_en": ocr_en,
        "llm_result": llm_result,
        "processing_time": {
            "ocr": ocr_time,
            "translation": translation_time,
            "llm": llm_time,
            "total": total_time
        }
    }
    
    # Enregistrer le résultat
    output_file = f"complete_test_{os.path.basename(image_path)}.json"
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print(f"Résultat enregistré dans: {output_file}")
    
    return result

def test_directory(dir_path, verbose=False, only_ocr=False, skip_llm=False):
    """
    Teste toutes les images d'un répertoire
    """
    if not os.path.exists(dir_path) or not os.path.isdir(dir_path):
        print(f"Le répertoire {dir_path} n'existe pas")
        return
    
    # Extensions d'images à tester
    extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.tif', '.webp']
    
    # Trouver toutes les images
    image_files = []
    for root, _, files in os.walk(dir_path):
        for file in files:
            if any(file.lower().endswith(ext) for ext in extensions):
                image_files.append(os.path.join(root, file))
    
    print(f"=== TEST DE {len(image_files)} IMAGES DANS {dir_path} ===")
    
    results = []
    for i, image_path in enumerate(image_files, 1):
        print(f"\nImage {i}/{len(image_files)}: {os.path.basename(image_path)}")
        result = test_image(image_path, verbose, only_ocr, skip_llm)
        if result:
            results.append(result)
    
    # Enregistrer les résultats combinés
    if results:
        output_file = f"batch_test_results.json"
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"\nRésultats combinés enregistrés dans: {output_file}")
    
    return results

def main():
    args = parse_args()
    
    # Test d'un répertoire
    if args.test_dir:
        test_directory(args.test_dir, args.verbose, args.only_ocr, args.skip_llm)
    # Test d'une seule image
    else:
        test_image(args.image_path, args.verbose, args.only_ocr, args.skip_llm)
    
    return 0

if __name__ == "__main__":
    sys.exit(main())