llm_ticket3/utils/image_extractor/extract_all_images.py

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

"""
Module combinant l'extracteur d'images standard avec l'extracteur d'images HTML.
Permet de récupérer toutes les images liées à un ticket, qu'elles soient directement
attachées ou référencées dans le HTML des messages.
"""

import os
import logging
from typing import Dict, List, Any, Set
from .html_image_extractor import extract_images_from_ticket

def enhance_ticket_data_with_images(ticket_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Récupère toutes les images associées à un ticket, en combinant l'extraction standard
    et l'extraction des images intégrées dans le HTML.
    
    Args:
        ticket_data: Données du ticket (avec la clé 'ticket_dir' contenant le chemin du répertoire)
        
    Returns:
        Données du ticket enrichies avec les chemins des images
    """
    result = ticket_data.copy()
    
    # Vérifier que le répertoire du ticket est spécifié
    ticket_dir = ticket_data.get("ticket_dir")
    if not ticket_dir or not os.path.exists(ticket_dir):
        logging.error("Répertoire du ticket non spécifié ou introuvable")
        result["images"] = []
        result["standard_images"] = []
        result["embedded_images"] = []
        result["images_count"] = 0
        result["images_stats"] = {"total": 0, "standard": 0, "embedded": 0}
        return result
    
    # 1. Récupérer les images standard depuis les pièces jointes
    standard_images = get_standard_images(ticket_data)
    
    # 2. Récupérer les images intégrées dans le HTML
    try:
        embedded_images = extract_images_from_ticket(ticket_dir)
    except Exception as e:
        logging.error(f"Erreur lors de l'extraction des images HTML: {e}")
        embedded_images = []
    
    # 3. Combiner les deux ensembles d'images et éliminer les doublons
    all_images = deduplicate_images(standard_images, embedded_images)
    
    # 4. Enrichir les données du ticket
    result["images"] = all_images
    result["standard_images"] = standard_images
    result["embedded_images"] = embedded_images
    result["images_count"] = len(all_images)
    result["images_stats"] = {
        "total": len(all_images),
        "standard": len(standard_images),
        "embedded": len(embedded_images)
    }
    
    return result

def get_standard_images(ticket_data: Dict[str, Any]) -> List[str]:
    """
    Récupère les images standard depuis les pièces jointes du ticket.
    
    Args:
        ticket_data: Données du ticket
        
    Returns:
        Liste des chemins des images standard
    """
    image_paths = []
    attachments = ticket_data.get("attachments", [])
    
    if not attachments:
        return []
    
    for attachment in attachments:
        # Vérifier si la pièce jointe est une image
        mimetype = attachment.get("mimetype", "").lower()
        name = attachment.get("name", "").lower()
        
        # Vérifier si c'est une image par mimetype ou extension
        is_image = (
            mimetype.startswith("image/") or 
            any(name.endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.svg'])
        )
        
        if is_image:
            # Utiliser le chemin local s'il existe
            path = attachment.get("local_path")
            if path and os.path.exists(path):
                image_paths.append(path)
    
    return image_paths

def deduplicate_images(standard_images: List[str], embedded_images: List[str]) -> List[str]:
    """
    Combine deux listes d'images en éliminant les doublons.
    
    Args:
        standard_images: Liste des chemins d'images standard
        embedded_images: Liste des chemins d'images intégrées
        
    Returns:
        Liste combinée des chemins d'images sans doublons
    """
    # Utiliser un ensemble pour éliminer les doublons
    unique_images = set(standard_images)
    
    # Ajouter les images intégrées qui ne sont pas déjà dans les images standard
    for img in embedded_images:
        unique_images.add(img)
    
    # Convertir en liste et vérifier que les chemins existent
    result = [img for img in unique_images if os.path.exists(img)]
    
    return result

if __name__ == "__main__":
    # Test simple
    import sys
    
    if len(sys.argv) > 1:
        ticket_dir = sys.argv[1]
    else:
        ticket_dir = "./output/ticket_T0241/T0241_20250409_141018"
    
    if not os.path.exists(ticket_dir):
        print(f"Répertoire introuvable: {ticket_dir}")
        sys.exit(1)
    
    # Charger les pièces jointes
    attachments_path = os.path.join(ticket_dir, "attachments_info.json")
    attachments = []
    
    if os.path.exists(attachments_path):
        import json
        with open(attachments_path, 'r', encoding='utf-8') as f:
            attachments = json.load(f)
    
    # Créer les données du ticket
    ticket_data = {
        "ticket_dir": ticket_dir,
        "attachments": attachments
    }
    
    # Enrichir avec les images
    result = enhance_ticket_data_with_images(ticket_data)
    
    # Afficher les résultats
    print(f"Images trouvées: {result['images_count']}")
    print(f"  Images standard: {len(result['standard_images'])}")
    print(f"  Images intégrées: {len(result['embedded_images'])}")
    
    if result["images"]:
        print("\nListe des images:")
        for i, img in enumerate(result["images"]):
            print(f"  {i+1}. {img}")