rpa_vision_v3/examples/test_ollama_integration.py

"""
Script de test pour l'intégration Ollama avec UIDetector

Ce script montre comment utiliser le UIDetector avec Ollama pour
détecter et classifier des éléments UI dans des screenshots.
"""

import sys
from pathlib import Path

# Ajouter le répertoire parent au path
sys.path.insert(0, str(Path(__file__).parent.parent))

from core.detection import (
    UIDetector,
    DetectionConfig,
    OllamaClient,
    check_ollama_available
)


def test_ollama_connection():
    """Tester la connexion à Ollama"""
    print("=" * 60)
    print("Test 1: Vérification de la connexion Ollama")
    print("=" * 60)
    
    if check_ollama_available():
        print("✅ Ollama est disponible!")
        
        # Lister les modèles
        client = OllamaClient()
        models = client.list_models()
        print(f"\nModèles disponibles: {models}")
        
        return True
    else:
        print("❌ Ollama n'est pas disponible")
        print("\nPour installer Ollama:")
        print("1. Visitez: https://ollama.ai")
        print("2. Téléchargez et installez Ollama")
        print("3. Lancez: ollama pull qwen2.5-vl")
        print("4. Vérifiez: ollama list")
        
        return False


def test_ui_detector_with_ollama(screenshot_path: str):
    """Tester le UIDetector avec Ollama"""
    print("\n" + "=" * 60)
    print("Test 2: Détection UI avec Ollama")
    print("=" * 60)
    
    # Créer client Ollama
    ollama_client = OllamaClient(model="qwen3-vl:8b")
    
    # Créer UIDetector
    config = DetectionConfig(
        vlm_model="qwen3-vl:8b",
        confidence_threshold=0.7,
        detect_regions=True
    )
    detector = UIDetector(config)
    
    # Brancher Ollama au détecteur
    detector.set_vlm_client(ollama_client)
    
    # Détecter éléments UI
    print(f"\nAnalyse du screenshot: {screenshot_path}")
    elements = detector.detect(screenshot_path)
    
    print(f"\n✅ {len(elements)} éléments UI détectés:")
    for i, elem in enumerate(elements, 1):
        print(f"\n  {i}. {elem.type.upper()} - {elem.role}")
        print(f"     Position: {elem.bbox}")
        print(f"     Label: {elem.label}")
        print(f"     Confiance: {elem.confidence:.2f}")


def test_element_classification():
    """Tester la classification d'éléments"""
    print("\n" + "=" * 60)
    print("Test 3: Classification d'éléments")
    print("=" * 60)
    
    client = OllamaClient(model="qwen2.5-vl")
    
    # Test avec une image fictive
    from PIL import Image
    import numpy as np
    
    # Créer une image de test (bouton bleu)
    img_array = np.zeros((50, 150, 3), dtype=np.uint8)
    img_array[:, :] = [0, 100, 200]  # Bleu
    test_image = Image.fromarray(img_array)
    
    # Classifier le type
    print("\nClassification du type...")
    type_result = client.classify_element_type(test_image)
    if type_result["success"]:
        print(f"✅ Type: {type_result['type']} (confiance: {type_result['confidence']:.2f})")
    else:
        print("❌ Échec de classification")
    
    # Classifier le rôle
    print("\nClassification du rôle...")
    role_result = client.classify_element_role(test_image, "button")
    if role_result["success"]:
        print(f"✅ Rôle: {role_result['role']} (confiance: {role_result['confidence']:.2f})")
    else:
        print("❌ Échec de classification")


def main():
    """Fonction principale"""
    print("\n" + "=" * 60)
    print("TEST D'INTÉGRATION OLLAMA + UIDetector")
    print("=" * 60)
    
    # Test 1: Connexion Ollama
    if not test_ollama_connection():
        print("\n⚠️  Ollama n'est pas disponible. Tests limités.")
        return
    
    # Test 2: Détection UI (si screenshot fourni)
    if len(sys.argv) > 1:
        screenshot_path = sys.argv[1]
        if Path(screenshot_path).exists():
            test_ui_detector_with_ollama(screenshot_path)
        else:
            print(f"\n❌ Screenshot non trouvé: {screenshot_path}")
    else:
        print("\n💡 Pour tester la détection UI:")
        print("   python test_ollama_integration.py <chemin_screenshot>")
    
    # Test 3: Classification
    test_element_classification()
    
    print("\n" + "=" * 60)
    print("Tests terminés!")
    print("=" * 60)


if __name__ == "__main__":
    main()