Update handler.py

handler.py

@@ -444,483 +444,6 @@ class EndpointHandler:
         except:
             pass
 
-# Test de fonctionnement si exécuté directement
-if __name__ == "__main__":
-    print("🧪 TEST DU HANDLER VERIFAI V2 FIXED")
-    print("=" * 50)
-    
-    try:
-        # Initialisation
-        handler = EndpointHandler()
-        
-        if handler.model is not None:
-            print("✅ Initialisation réussie")
-            
-            # Test avec une image simple
-            print("🔄 Test avec image de base...")
-            test_img = Image.new('RGB', (224, 224), color='red')
-            buffer = io.BytesIO()
-            test_img.save(buffer, format='JPEG')
-            test_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
-            
-            result = handler({"inputs": test_data})
-            print(f"📊 Résultat: {result['status']}")
-            if result['status'] == 'success':
-                print(f"🎯 Prédiction: {result['predicted_class_name']} ({result['confidence']:.3f})")
-                print("✅ Handler fonctionnel!")
-            else:
-                print(f"❌ Erreur: {result.get('error', 'Inconnue')}")
-        else:
-            print("❌ Échec de l'initialisation")
-            
-    except Exception as e:
-        print(f"❌ Erreur de test: {e}")
-        print(f"🔍 Traceback: {traceback.format_exc()}") from typing import Any, Dict
-import torch
-from torchvision import transforms
-from PIL import Image, ImageDraw, ImageOps
-import base64
-import io
-import numpy as np
-from transformers import AutoModelForImageClassification, AutoImageProcessor, AutoConfig
-import torch.nn.functional as F
-import json
-import re
-import gc
-import sys
-import traceback
-
-# Gestion des dépendances optionnelles
-HAS_MATPLOTLIB = False
-try:
-    import matplotlib.pyplot as plt
-    import matplotlib.cm as cm
-    HAS_MATPLOTLIB = True
-    print("✅ Matplotlib disponible - Grad-CAM avancé activé")
-except ImportError:
-    print("⚠️  Matplotlib non disponible - Utilisation de PIL pour Grad-CAM")
-
-class OptimizedGradCAM:
-    """Version optimisée de Grad-CAM avec nettoyage automatique"""
-    
-    def __init__(self, model, target_layer):
-        self.model = model
-        self.target_layer = target_layer
-        self.gradients = None
-        self.activations = None
-        self.hooks = []
-        
-        # Enregistrer les hooks avec nettoyage automatique
-        if target_layer is not None:
-            hook1 = self.target_layer.register_backward_hook(self.save_gradients)
-            hook2 = self.target_layer.register_forward_hook(self.save_activations)
-            self.hooks = [hook1, hook2]
-        else:
-            print("⚠️  Aucune couche cible trouvée - Grad-CAM désactivé")
-    
-    def save_gradients(self, module, grad_input, grad_output):
-        if grad_output[0] is not None:
-            self.gradients = grad_output[0].detach()
-    
-    def save_activations(self, module, input, output):
-        self.activations = output.detach()
-    
-    def generate_cam(self, input_tensor, class_idx=None):
-        """Génère la carte de saillance Grad-CAM"""
-        if self.target_layer is None:
-            return None
-            
-        try:
-            # Forward pass
-            output = self.model(input_tensor)
-            
-            if class_idx is None:
-                class_idx = output.logits.argmax(dim=1).item()
-            
-            # Backward pass
-            self.model.zero_grad()
-            output.logits[0, class_idx].backward(retain_graph=False)
-            
-            if self.gradients is None or self.activations is None:
-                print("⚠️  Gradients ou activations manquants")
-                return None
-            
-            # Generate CAM
-            gradients = self.gradients[0]  # (C, H, W)
-            activations = self.activations[0]  # (C, H, W)
-            
-            # Moyenne globale des gradients
-            weights = torch.mean(gradients, dim=(1, 2))  # (C,)
-            
-            # CAM = somme pondérée des activations
-            cam = torch.zeros(activations.shape[1:], device=activations.device)  # (H, W)
-            for i, w in enumerate(weights):
-                cam += w * activations[i, :, :]
-            
-            # ReLU et normalisation
-            cam = F.relu(cam)
-            if cam.max() > 0:
-                cam = cam / cam.max()
-            
-            return cam.detach().cpu().numpy()
-            
-        except Exception as e:
-            print(f"⚠️  Erreur lors de la génération CAM: {e}")
-            return None
-        finally:
-            # Nettoyage explicite
-            if self.gradients is not None:
-                self.gradients = None
-            if self.activations is not None:
-                self.activations = None
-    
-    def cleanup(self):
-        """Nettoie les hooks et libère la mémoire"""
-        for hook in self.hooks:
-            try:
-                hook.remove()
-            except:
-                pass
-        self.hooks = []
-        self.gradients = None
-        self.activations = None
-    
-    def __del__(self):
-        """Nettoyage automatique lors de la destruction"""
-        self.cleanup()
-
-def get_last_conv_layer_safe(model):
-    """Trouve la dernière couche de convolution de manière sécurisée"""
-    try:
-        last_conv = None
-        conv_layers = []
-        
-        for name, module in model.named_modules():
-            if isinstance(module, (torch.nn.Conv2d, torch.nn.AdaptiveAvgPool2d)):
-                conv_layers.append((name, module))
-        
-        if conv_layers:
-            last_conv = conv_layers[-1][1]
-            print(f"✅ Couche cible trouvée: {conv_layers[-1][0]}")
-        else:
-            print("⚠️  Aucune couche de convolution trouvée")
-            
-        return last_conv
-    except Exception as e:
-        print(f"⚠️  Erreur lors de la recherche de couche: {e}")
-        return None
-
-def create_gradcam_overlay_pil(original_image, cam_array):
-    """Crée une superposition Grad-CAM en utilisant PIL (sans matplotlib)"""
-    try:
-        if cam_array is None:
-            return None
-            
-        # Convertir CAM en image
-        cam_normalized = (cam_array * 255).astype(np.uint8)
-        cam_img = Image.fromarray(cam_normalized, mode='L')
-        
-        # Redimensionner au format de l'image originale
-        cam_resized = cam_img.resize(original_image.size, Image.Resampling.LANCZOS)
-        
-        # Créer une heatmap colorée (rouge pour les zones importantes)
-        # Convertir en RGB et appliquer une colormap simple
-        cam_array_resized = np.array(cam_resized)
-        
-        # Créer une colormap simple (bleu -> rouge)
-        heatmap = np.zeros((cam_array_resized.shape[0], cam_array_resized.shape[1], 3), dtype=np.uint8)
-        heatmap[:, :, 0] = cam_array_resized  # Rouge
-        heatmap[:, :, 2] = 255 - cam_array_resized  # Bleu inversé
-        
-        heatmap_img = Image.fromarray(heatmap, 'RGB')
-        
-        # Mélanger avec l'image originale
-        blended = Image.blend(original_image.convert('RGB'), heatmap_img, alpha=0.4)
-        
-        # Convertir en base64
-        buffer = io.BytesIO()
-        blended.save(buffer, format='PNG', optimize=True)
-        buffer.seek(0)
-        
-        return base64.b64encode(buffer.getvalue()).decode('utf-8')
-        
-    except Exception as e:
-        print(f"⚠️  Erreur lors de la création de l'overlay PIL: {e}")
-        return None
-
-def create_gradcam_overlay_matplotlib(original_image, cam_array):
-    """Crée une superposition Grad-CAM en utilisant matplotlib (si disponible)"""
-    try:
-        if not HAS_MATPLOTLIB or cam_array is None:
-            return None
-            
-        # Redimensionner CAM
-        cam_resized = np.array(Image.fromarray((cam_array * 255).astype(np.uint8)).resize(
-            original_image.size, Image.Resampling.LANCZOS
-        )) / 255.0
-        
-        # Créer la figure
-        fig, ax = plt.subplots(figsize=(8, 8), dpi=100)
-        ax.imshow(original_image)
-        ax.imshow(cam_resized, cmap='jet', alpha=0.5)
-        ax.axis('off')
-        
-        # Sauvegarder en base64
-        buffer = io.BytesIO()
-        plt.savefig(buffer, format='png', bbox_inches='tight', pad_inches=0, dpi=100)
-        plt.close(fig)  # Important: fermer la figure
-        buffer.seek(0)
-        
-        return base64.b64encode(buffer.getvalue()).decode('utf-8')
-        
-    except Exception as e:
-        print(f"⚠️  Erreur lors de la création de l'overlay matplotlib: {e}")
-        if 'fig' in locals():
-            plt.close(fig)
-        return None
-
-class EndpointHandler:
-    def __init__(self, path=""):
-        print("🚀 VerifAI Handler V2 FIXED - Initialisation")
-        print("📋 Modèle: haywoodsloan/ai-image-detector-deploy (Version Corrigée)")
-        
-        self.model = None
-        self.processor = None
-        self.grad_cam = None
-        self.model_labels = {}
-        
-        try:
-            # Vérification de la disponibilité du modèle
-            self.model_name = "haywoodsloan/ai-image-detector-deploy"
-            
-            if not self._verify_model_exists():
-                raise Exception(f"Modèle {self.model_name} non accessible")
-            
-            # Chargement du modèle avec gestion d'erreurs
-            print("🔄 Chargement du modèle...")
-            self.processor = AutoImageProcessor.from_pretrained(self.model_name)
-            self.model = AutoModelForImageClassification.from_pretrained(
-                self.model_name,
-                torch_dtype=torch.float32  # Force float32 pour la compatibilité
-            )
-            self.model.eval()
-            
-            # Configuration Grad-CAM sécurisée
-            target_layer = get_last_conv_layer_safe(self.model)
-            if target_layer is not None:
-                self.grad_cam = OptimizedGradCAM(self.model, target_layer)
-                print("✅ Grad-CAM activé")
-            else:
-                print("⚠️  Grad-CAM désactivé (aucune couche compatible)")
-            
-            # Récupérer les labels
-            if hasattr(self.model.config, 'id2label'):
-                self.model_labels = self.model.config.id2label
-            else:
-                self.model_labels = {0: "Real", 1: "Fake"}  # Fallback
-            
-            print("✅ Modèle chargé avec succès")
-            print(f"📋 Étiquettes du modèle: {self.model_labels}")
-            print("🎯 VerifAI Handler V2 FIXED prêt!")
-            
-        except Exception as e:
-            print(f"❌ Erreur lors de l'initialisation: {e}")
-            print(f"🔍 Traceback: {traceback.format_exc()}")
-            # Ne pas faire échouer l'initialisation, mais signaler l'erreur
-            self.model = None
-            self.processor = None
-    
-    def _verify_model_exists(self):
-        """Vérifie que le modèle existe avant de le charger"""
-        try:
-            config = AutoConfig.from_pretrained(self.model_name)
-            print(f"✅ Modèle {self.model_name} vérifié")
-            return True
-        except Exception as e:
-            print(f"❌ Modèle {self.model_name} non accessible: {e}")
-            return False
-    
-    def _normalize_label(self, label: str) -> str:
-        """Normalise les étiquettes pour qu'elles soient cohérentes."""
-        if not isinstance(label, str):
-            label = str(label)
-            
-        label_lower = label.lower()
-        if re.search(r'real|human|authentic', label_lower):
-            return "Human"
-        if re.search(r'fake|generated|ai|artificial', label_lower):
-            return "AI Generated"
-        return "Unknown"
-    
-    def _cleanup_memory(self):
-        """Nettoie la mémoire explicitement"""
-        try:
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            gc.collect()
-        except:
-            pass
-    
-    def __call__(self, data):
-        # Vérification de l'état du handler
-        if self.model is None or self.processor is None:
-            return {
-                "status": "error",
-                "error": "Handler non initialisé correctement",
-                "prediction": 0,
-                "predicted_class_name": "Error",
-                "confidence": 0.0,
-                "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
-                "cam_image": None,
-                "version": "2.0-fixed",
-                "handler_name": "VerifAI Handler V2 FIXED"
-            }
-        
-        try:
-            # Traitement de l'image avec validation
-            image_data = data.get("inputs") or data
-            if not image_data:
-                raise ValueError("Aucune donnée d'image fournie")
-            
-            # Décodage sécurisé de l'image
-            try:
-                image_bytes = base64.b64decode(image_data)
-                image = Image.open(io.BytesIO(image_bytes))
-                
-                # Validation et conversion
-                if image.mode != 'RGB':
-                    image = image.convert('RGB')
-                    
-                # Validation de la taille
-                if image.size[0] * image.size[1] > 4096 * 4096:
-                    image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
-                    print("⚠️  Image redimensionnée pour éviter les problèmes de mémoire")
-                    
-            except Exception as e:
-                raise ValueError(f"Erreur lors du décodage de l'image: {e}")
-
-            # Prédiction avec gestion d'erreurs
-            print("🔄 VerifAI V2 FIXED - Analyse en cours...")
-            
-            try:
-                inputs = self.processor(image, return_tensors="pt")
-                
-                with torch.no_grad():
-                    outputs = self.model(**inputs)
-                    logits = outputs.logits
-                    probabilities = F.softmax(logits, dim=-1)[0]
-                    predicted_class_id = logits.argmax().item()
-                    
-            except Exception as e:
-                raise RuntimeError(f"Erreur lors de l'inférence: {e}")
-            
-            # Traitement des résultats
-            class_probs = {}
-            for class_id, prob in enumerate(probabilities):
-                label_str = self.model_labels.get(class_id, f"Class {class_id}")
-                normalized_label = self._normalize_label(label_str)
-                if normalized_label != "Unknown":
-                    class_probs[normalized_label] = float(prob)
-
-            # S'assurer que les deux classes existent
-            class_probs.setdefault("Human", 0.0)
-            class_probs.setdefault("AI Generated", 0.0)
-
-            prediction_label = self._normalize_label(self.model_labels.get(predicted_class_id, "Unknown"))
-            confidence = class_probs.get(prediction_label, 0.0)
-            
-            # Déterminer l'ID de prédiction pour la compatibilité
-            prediction_id = 1 if prediction_label == "AI Generated" else 0
-            
-            print(f"🔍 VerifAI V2 FIXED Résultat: {prediction_label} (confiance: {confidence:.3f})")
-
-            # Génération du Grad-CAM avec fallback
-            cam_image_b64 = None
-            if self.grad_cam is not None:
-                try:
-                    print("🎨 Génération du Grad-CAM...")
-                    cam = self.grad_cam.generate_cam(inputs['pixel_values'], predicted_class_id)
-                    
-                    if cam is not None:
-                        # Essayer matplotlib d'abord, puis PIL
-                        if HAS_MATPLOTLIB:
-                            cam_image_b64 = create_gradcam_overlay_matplotlib(image, cam)
-                        
-                        if cam_image_b64 is None:
-                            cam_image_b64 = create_gradcam_overlay_pil(image, cam)
-                        
-                        if cam_image_b64:
-                            print("✅ Grad-CAM généré avec succès")
-                        else:
-                            print("⚠️  Échec de la génération Grad-CAM")
-                    
-                except Exception as e:
-                    print(f"⚠️  Erreur Grad-CAM: {e}")
-                    cam_image_b64 = None
-
-            # Nettoyage mémoire
-            self._cleanup_memory()
-
-            # Construction de la réponse compatible
-            return {
-                "status": "success",
-                "prediction": prediction_id,
-                "predicted_class_name": prediction_label,
-                "confidence": confidence,
-                "class_probabilities": class_probs,
-                "cam_image": cam_image_b64,
-                "model_info": {
-                    "model_name": self.model_name,
-                    "handler_version": "verifai-v2-fixed",
-                    "precision_mode": "high",
-                    "raw_prediction_id": predicted_class_id,
-                    "raw_labels": self.model_labels,
-                    "grad_cam_method": "matplotlib" if HAS_MATPLOTLIB else "pil"
-                },
-                "reliability": "TRÈS ÉLEVÉE",
-                "version": "2.0-fixed",
-                "handler_name": "VerifAI Handler V2 FIXED",
-                "deployment_note": "VERIFAI HANDLER V2 FIXED - PRODUCTION READY",
-                "fixes_applied": [
-                    "Gestion d'erreurs robuste",
-                    "Fallback PIL pour Grad-CAM",
-                    "Nettoyage mémoire automatique",
-                    "Validation d'entrée renforcée"
-                ]
-            }
-
-        except Exception as e:
-            print(f"❌ Erreur dans VerifAI Handler V2 FIXED: {e}")
-            print(f"🔍 Traceback: {traceback.format_exc()}")
-            
-            # Nettoyage en cas d'erreur
-            self._cleanup_memory()
-            
-            return {
-                "status": "error",
-                "error": str(e),
-                "prediction": 0,
-                "predicted_class_name": "Error",
-                "confidence": 0.0,
-                "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
-                "cam_image": None,
-                "version": "2.0-fixed",
-                "handler_name": "VerifAI Handler V2 FIXED",
-                "error_details": {
-                    "error_type": type(e).__name__,
-                    "traceback": traceback.format_exc()[-500:]  # Dernières 500 chars
-                }
-            }
-    
-    def __del__(self):
-        """Nettoyage lors de la destruction de l'instance"""
-        try:
-            if hasattr(self, 'grad_cam') and self.grad_cam is not None:
-                self.grad_cam.cleanup()
-            self._cleanup_memory()
-        except:
-            pass
-
 # Test de fonctionnement si exécuté directement
 if __name__ == "__main__":
     print("🧪 TEST DU HANDLER VERIFAI V2 FIXED")