from typing import Any, Dict import torch from torchvision import transforms from PIL import Image import base64 import io import numpy as np from transformers import AutoModelForImageClassification, AutoImageProcessor from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.image import show_cam_on_image import warnings warnings.filterwarnings("ignore") class EndpointHandler: def __init__(self, model_dir: str = "haywoodsloan/ai-image-detector-deploy", **kwargs: Any): """ Initialise le handler avec le modèle haywoodsloan/ai-image-detector-deploy et configure Grad-CAM pour les cartes de saillance. """ # Forcer l'utilisation du modèle Hugging Face Hub même lors du déploiement model_name = "haywoodsloan/ai-image-detector-deploy" print(f"Initialisation du handler avec le modèle : {model_name}") print(f"Répertoire de déploiement : {model_dir}") # Charger le modèle et le processeur depuis Hugging Face Hub self.model = AutoModelForImageClassification.from_pretrained(model_name) self.model.eval() self.processor = AutoImageProcessor.from_pretrained(model_name) # Configuration de Grad-CAM - Adaptation pour différentes architectures # Le modèle haywoodsloan/ai-image-detector-deploy est généralement basé sur ViT ou CNN self.target_layer = self._find_target_layer() self.cam = GradCAM( model=self.model, target_layers=[self.target_layer] ) # Mapping des classes pour le détecteur d'IA # Configuration basée sur la structure du modèle haywoodsloan self.class_names = { 0: "Image Réelle", 1: "Image Générée par IA" } # Seuils de confiance pour l'interprétation self.confidence_thresholds = { "très_élevée": 0.9, "élevée": 0.75, "moyenne": 0.6, "faible": 0.4 } print("Handler initialisé avec succès!") def _find_target_layer(self): """ Trouve automatiquement la couche cible appropriée pour Grad-CAM selon l'architecture du modèle. """ try: # Pour les modèles Vision Transformer (ViT) if hasattr(self.model, 'vit'): if hasattr(self.model.vit, 'encoder'): return self.model.vit.encoder.layer[-1].layernorm_before elif hasattr(self.model.vit, 'layers'): return self.model.vit.layers[-1].norm1 # Pour les modèles Swin Transformer elif hasattr(self.model, 'swin'): return self.model.swin.encoder.layers[-1].blocks[-1].layernorm_before # Pour les modèles avec backbone elif hasattr(self.model, 'backbone'): if hasattr(self.model.backbone, 'layers'): return self.model.backbone.layers[-1].blocks[-1].norm1 else: # Fallback pour backbone CNN return list(self.model.backbone.children())[-2] # Pour les modèles ConvNeXt elif hasattr(self.model, 'convnext'): return self.model.convnext.encoder.stages[-1].layers[-1].layernorm # Pour les modèles ResNet ou autres architectures CNN elif hasattr(self.model, 'resnet'): return self.model.resnet.layer4[-1].bn2 # Fallback générique - chercher la dernière couche de normalisation else: # Parcourir tous les modules pour trouver une couche appropriée modules = list(self.model.named_modules()) for name, module in reversed(modules): if any(layer_type in name.lower() for layer_type in ['layernorm', 'batchnorm', 'norm']): if 'classifier' not in name.lower(): print(f"Couche cible trouvée : {name}") return module # Si aucune couche de normalisation trouvée, utiliser l'avant-dernière couche children = list(self.model.children()) if len(children) > 1: return children[-2] else: return children[-1] except Exception as e: print(f"Erreur lors de la recherche de la couche cible: {e}") # Fallback final - utiliser la première couche trouvée children = list(self.model.children()) return children[-2] if len(children) > 1 else children[0] def _interpret_confidence(self, confidence: float, predicted_class: str) -> str: """ Interprète le niveau de confiance et génère un message explicatif. """ if confidence >= self.confidence_thresholds["très_élevée"]: level = "très élevée" reliability = "Très fiable" elif confidence >= self.confidence_thresholds["élevée"]: level = "élevée" reliability = "Fiable" elif confidence >= self.confidence_thresholds["moyenne"]: level = "moyenne" reliability = "Moyennement fiable" else: level = "faible" reliability = "Peu fiable" interpretation = f"Confiance {level} ({confidence:.1%}) - {reliability}. " if predicted_class == "Image Générée par IA": if confidence >= 0.8: interpretation += "L'image présente des caractéristiques typiques d'une génération par IA." elif confidence >= 0.6: interpretation += "L'image pourrait être générée par IA, mais nécessite une vérification supplémentaire." else: interpretation += "Classification incertaine - analyse manuelle recommandée." else: if confidence >= 0.8: interpretation += "L'image semble authentique avec des caractéristiques naturelles." elif confidence >= 0.6: interpretation += "L'image semble réelle, mais avec quelques éléments à vérifier." else: interpretation += "Classification incertaine - analyse manuelle recommandée." return interpretation def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: """ Traite une image et retourne la prédiction avec la carte de saillance. Args: data: Dictionnaire contenant l'image encodée en base64 Returns: Dictionnaire avec la prédiction, confiance et carte de saillance """ try: print("Début du traitement de l'image...") # Décoder l'image depuis une chaîne base64 if isinstance(data["inputs"], str): image_data = base64.b64decode(data["inputs"]) else: # Si c'est déjà des bytes image_data = data["inputs"] image = Image.open(io.BytesIO(image_data)).convert("RGB") print(f"Image chargée avec succès : {image.size}") # Préprocesser l'image pour le modèle inputs = self.processor(images=image, return_tensors="pt") input_tensor = inputs["pixel_values"] print("Génération de la carte de saillance Grad-CAM...") # Générer la carte de saillance avec Grad-CAM try: # Correction spécifique pour Swin Transformer v2 # Créer une classe wrapper pour le modèle import torch.nn as nn class ModelWrapper(nn.Module): def __init__(self, model): super().__init__() self.model = model def forward(self, x): outputs = self.model(x) # Extraire les logits des outputs de Swin v2 return outputs.logits # Créer le wrapper wrapped_model = ModelWrapper(self.model) # Créer un nouveau GradCAM avec le modèle wrapper from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget wrapped_cam = GradCAM( model=wrapped_model, target_layers=[self.target_layer] ) # Obtenir la classe prédite with torch.no_grad(): outputs = self.model(input_tensor) predicted_class_idx = torch.argmax(outputs.logits, dim=1).item() # Créer le target pour GradCAM targets = [ClassifierOutputTarget(predicted_class_idx)] # Générer la carte de saillance grayscale_cam = wrapped_cam(input_tensor=input_tensor, targets=targets)[0] # Redimensionner l'image originale pour correspondre à la carte de saillance cam_height, cam_width = grayscale_cam.shape image_resized = image.resize((cam_width, cam_height)) image_np = np.array(image_resized).astype(np.float32) / 255.0 # Superposer la carte de chaleur sur l'image visualization = show_cam_on_image(image_np, grayscale_cam, use_rgb=True) # Convertir la carte de chaleur en base64 buffered = io.BytesIO() Image.fromarray(visualization).save(buffered, format="PNG") cam_image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8") except Exception as e: print(f"Erreur lors de la génération de Grad-CAM: {e}") cam_image_base64 = None print("Exécution de la prédiction...") # Obtenir la prédiction du modèle with torch.no_grad(): outputs = self.model(**inputs) logits = outputs.logits probabilities = torch.nn.functional.softmax(logits, dim=1) predicted_class = torch.argmax(probabilities, dim=1).item() confidence = probabilities[0][predicted_class].item() # Calculer le score de confiance pour chaque classe class_probabilities = {} for i, prob in enumerate(probabilities[0].tolist()): class_name = self.class_names.get(i, f"Classe {i}") class_probabilities[class_name] = round(prob, 4) # Générer l'interprétation predicted_class_name = self.class_names.get(predicted_class, f"Classe {predicted_class}") interpretation = self._interpret_confidence(confidence, predicted_class_name) # Score de détection d'IA (probabilité que ce soit une IA) ai_detection_score = probabilities[0][1].item() if len(probabilities[0]) > 1 else 0.0 result = { "prediction": predicted_class, "predicted_class_name": predicted_class_name, "confidence": round(confidence, 4), "ai_detection_score": round(ai_detection_score, 4), "class_probabilities": class_probabilities, "interpretation": interpretation, "status": "success", "model_used": "haywoodsloan/ai-image-detector-deploy" } # Ajouter la carte Grad-CAM si disponible if cam_image_base64: result["cam_image"] = cam_image_base64 result["grad_cam_available"] = True else: result["grad_cam_available"] = False result["grad_cam_error"] = "Impossible de générer la carte de saillance" print(f"Traitement terminé avec succès! Prédiction: {predicted_class_name}, Confiance: {confidence:.2%}") return result except Exception as e: print(f"Erreur lors du traitement: {e}") return { "error": str(e), "status": "error", "model_used": "haywoodsloan/ai-image-detector-deploy" } # Test local du handler (optionnel) if __name__ == "__main__": import os try: print("Test d'initialisation du handler...") handler = EndpointHandler() print("Handler initialisé avec succès!") # Test avec une image d'exemple test_image_path = "test_image.jpg" if os.path.exists(test_image_path): print(f"Test avec l'image : {test_image_path}") with open(test_image_path, "rb") as f: image_bytes = f.read() input_data = {"inputs": base64.b64encode(image_bytes).decode("utf-8")} output = handler(input_data) print("\n=== RÉSULTATS DU TEST ===") print(f"Statut: {output.get('status', 'N/A')}") print(f"Prédiction: {output.get('predicted_class_name', 'N/A')}") print(f"Confiance: {output.get('confidence', 0):.2%}") print(f"Score de détection IA: {output.get('ai_detection_score', 0):.2%}") print(f"Grad-CAM disponible: {output.get('grad_cam_available', False)}") print(f"Interprétation: {output.get('interpretation', 'N/A')}") if 'class_probabilities' in output: print("\nProbabilités par classe:") for class_name, prob in output['class_probabilities'].items(): print(f" {class_name}: {prob:.2%}") else: print(f"Aucune image de test trouvée : {test_image_path}") print("Placez une image de test dans le répertoire pour tester le handler.") print("Vous pouvez utiliser n'importe quel format d'image (JPG, PNG, etc.)") except Exception as e: print(f"Erreur lors de l'initialisation ou du test: {e}") import traceback traceback.print_exc()