Update handler.py

handler.py

@@ -1,282 +1,89 @@
-from typing import Any, Dict
 import torch
-from torchvision import transforms
-from PIL import Image, ImageDraw, ImageOps
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
 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
+import signal
+from contextlib import timeout as context_timeout
 
-# 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 TimeoutError(Exception):
+    pass
 
-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
+def timeout_handler(signum, frame):
+    raise TimeoutError("Opération timeout")
 
 class EndpointHandler:
     def __init__(self, path=""):
-        print("🚀 VerifAI Handler V2 FIXED - Initialisation")
-        print("📋 Modèle: haywoodsloan/ai-image-detector-deploy (Version Corrigée)")
+        print("🚀 VerifAI Handler V3 ULTRA LIGHT - Initialisation")
+        print("⚡ Version optimisée sans timeouts")
         
         self.model = None
         self.processor = None
-        self.grad_cam = None
         self.model_labels = {}
+        self.model_name = "haywoodsloan/ai-image-detector-deploy"
+        
+        # Timeout pour l'initialisation (30 secondes max)
+        signal.signal(signal.SIGALRM, timeout_handler)
+        signal.alarm(30)
         
         try:
-            # Vérification de la disponibilité du modèle
-            self.model_name = "haywoodsloan/ai-image-detector-deploy"
+            print("🔄 Chargement rapide du modèle...")
             
-            if not self._verify_model_exists():
-                raise Exception(f"Modèle {self.model_name} non accessible")
+            # Chargement simplifié et rapide
+            self.processor = AutoImageProcessor.from_pretrained(
+                self.model_name,
+                trust_remote_code=False
+            )
             
-            # 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é
+                torch_dtype=torch.float32,
+                trust_remote_code=False,
+                low_cpu_mem_usage=True  # Optimisation mémoire
             )
             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
+            # Labels simples
             if hasattr(self.model.config, 'id2label'):
                 self.model_labels = self.model.config.id2label
             else:
-                self.model_labels = {0: "Real", 1: "Fake"}  # Fallback
+                self.model_labels = {0: "Real", 1: "Fake"}
             
-            print("✅ Modèle chargé avec succès")
-            print(f"📋 Étiquettes du modèle: {self.model_labels}")
-            print("🎯 VerifAI Handler V2 FIXED prêt!")
+            print("✅ Modèle chargé (version ultra-light)")
+            print(f"📋 Labels: {self.model_labels}")
+            print("🎯 Handler 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
+        except TimeoutError:
+            print("❌ Timeout lors de l'initialisation")
             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
+            print(f"❌ Erreur initialisation: {e}")
+            self.model = None
+            self.processor = None
+        finally:
+            signal.alarm(0)  # Désactiver le timeout
     
     def _normalize_label(self, label: str) -> str:
-        """Normalise les étiquettes pour qu'elles soient cohérentes."""
+        """Normalise les labels rapidement"""
         if not isinstance(label, str):
             label = str(label)
-            
+        
         label_lower = label.lower()
-        if re.search(r'real|human|authentic', label_lower):
+        if any(word in label_lower for word in ['real', 'human', 'authentic']):
             return "Human"
-        if re.search(r'fake|generated|ai|artificial', label_lower):
+        if any(word in label_lower for word in ['fake', 'generated', 'ai', 'artificial']):
             return "AI Generated"
         return "Unknown"
     
     def _cleanup_memory(self):
-        """Nettoie la mémoire explicitement"""
+        """Nettoyage mémoire rapide"""
         try:
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
@@ -285,49 +92,60 @@ class EndpointHandler:
             pass
     
     def __call__(self, data):
-        # Vérification de l'état du handler
+        # Vérification rapide
         if self.model is None or self.processor is None:
             return {
                 "status": "error",
-                "error": "Handler non initialisé correctement",
+                "error": "Handler non initialisé",
                 "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"
+                "version": "3.0-ultra-light",
+                "handler_name": "VerifAI Handler V3 ULTRA LIGHT"
             }
         
+        # Timeout pour le traitement (15 secondes max)
+        signal.signal(signal.SIGALRM, timeout_handler)
+        signal.alarm(15)
+        
         try:
-            # Traitement de l'image avec validation
+            print("🔄 Traitement ultra-rapide...")
+            
+            # Extraction image avec validation minimale
             image_data = data.get("inputs") or data
             if not image_data:
-                raise ValueError("Aucune donnée d'image fournie")
+                raise ValueError("Pas de données image")
             
-            # Décodage sécurisé de l'image
+            # Décodage rapide
             try:
+                # Nettoyer le base64 si nécessaire
+                if isinstance(image_data, str) and image_data.startswith('data:'):
+                    image_data = image_data.split(',', 1)[1]
+                
                 image_bytes = base64.b64decode(image_data)
                 image = Image.open(io.BytesIO(image_bytes))
                 
-                # Validation et conversion
+                # Conversion rapide
                 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")
+                
+                # Redimensionnement agressif si trop grand
+                if image.size[0] * image.size[1] > 1024 * 1024:  # Plus de 1MP
+                    image = image.resize((512, 512), Image.Resampling.LANCZOS)
+                    print("⚠️  Image redimensionnée (optimisation)")
                     
             except Exception as e:
-                raise ValueError(f"Erreur lors du décodage de l'image: {e}")
+                raise ValueError(f"Erreur décodage image: {e}")
 
-            # Prédiction avec gestion d'erreurs
-            print("🔄 VerifAI V2 FIXED - Analyse en cours...")
+            print("🧠 Inférence rapide...")
             
+            # Traitement minimal et rapide
             try:
                 inputs = self.processor(image, return_tensors="pt")
                 
+                # Inférence avec timeout interne
                 with torch.no_grad():
                     outputs = self.model(**inputs)
                     logits = outputs.logits
@@ -335,90 +153,75 @@ class EndpointHandler:
                     predicted_class_id = logits.argmax().item()
                     
             except Exception as e:
-                raise RuntimeError(f"Erreur lors de l'inférence: {e}")
+                raise RuntimeError(f"Erreur inférence: {e}")
             
-            # Traitement des résultats
+            # Traitement des résultats - Version rapide
             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)
+                if class_id < len(self.model_labels):
+                    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
+            # Fallback pour assurer les deux classes
             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
+            print(f"🎯 Résultat: {prediction_label} ({confidence:.3f})")
 
-            # Nettoyage mémoire
+            # Nettoyage rapide
             self._cleanup_memory()
 
-            # Construction de la réponse compatible
+            # Réponse optimisée (sans Grad-CAM pour éviter les timeouts)
             return {
                 "status": "success",
                 "prediction": prediction_id,
                 "predicted_class_name": prediction_label,
                 "confidence": confidence,
                 "class_probabilities": class_probs,
-                "cam_image": cam_image_b64,
+                "cam_image": None,  # Désactivé temporairement
                 "model_info": {
                     "model_name": self.model_name,
-                    "handler_version": "verifai-v2-fixed",
-                    "precision_mode": "high",
+                    "handler_version": "verifai-v3-ultra-light",
+                    "precision_mode": "fast",
                     "raw_prediction_id": predicted_class_id,
                     "raw_labels": self.model_labels,
-                    "grad_cam_method": "matplotlib" if HAS_MATPLOTLIB else "pil"
+                    "grad_cam_method": "disabled_for_speed"
                 },
-                "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"
-                ]
+                "reliability": "ÉLEVÉE",
+                "version": "3.0-ultra-light",
+                "handler_name": "VerifAI Handler V3 ULTRA LIGHT",
+                "optimizations": [
+                    "Grad-CAM désactivé (évite timeouts)",
+                    "Chargement modèle optimisé",
+                    "Timeouts internes appliqués",
+                    "Nettoyage mémoire agressif"
+                ],
+                "performance_note": "Version rapide sans visualisation Grad-CAM"
             }
 
+        except TimeoutError:
+            print("❌ Timeout lors du traitement")
+            return {
+                "status": "error",
+                "error": "Timeout - traitement trop long",
+                "prediction": 0,
+                "predicted_class_name": "Error",
+                "confidence": 0.0,
+                "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
+                "cam_image": None,
+                "version": "3.0-ultra-light",
+                "handler_name": "VerifAI Handler V3 ULTRA LIGHT"
+            }
         except Exception as e:
-            print(f"❌ Erreur dans VerifAI Handler V2 FIXED: {e}")
+            print(f"❌ Erreur traitement: {e}")
             print(f"🔍 Traceback: {traceback.format_exc()}")
             
-            # Nettoyage en cas d'erreur
-            self._cleanup_memory()
-            
             return {
                 "status": "error",
                 "error": str(e),
@@ -427,52 +230,45 @@ class EndpointHandler:
                 "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",
+                "version": "3.0-ultra-light",
+                "handler_name": "VerifAI Handler V3 ULTRA LIGHT",
                 "error_details": {
                     "error_type": type(e).__name__,
-                    "traceback": traceback.format_exc()[-500:]  # Dernières 500 chars
+                    "traceback": traceback.format_exc()[-300:]
                 }
             }
-    
-    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()
+        finally:
+            signal.alarm(0)  # Désactiver le timeout
             self._cleanup_memory()
-        except:
-            pass
 
-# Test de fonctionnement si exécuté directement
+# Test de fonctionnement
 if __name__ == "__main__":
-    print("🧪 TEST DU HANDLER VERIFAI V2 FIXED")
+    print("🧪 TEST HANDLER V3 ULTRA LIGHT")
     print("=" * 50)
     
     try:
-        # Initialisation
         handler = EndpointHandler()
         
         if handler.model is not None:
-            print("✅ Initialisation réussie")
+            print("✅ Initialisation OK")
             
-            # Test avec une image simple
-            print("🔄 Test avec image de base...")
+            # Test rapide
             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')
             
+            print("🔄 Test inférence...")
             result = handler({"inputs": test_data})
-            print(f"📊 Résultat: {result['status']}")
+            
+            print(f"📊 Statut: {result['status']}")
             if result['status'] == 'success':
                 print(f"🎯 Prédiction: {result['predicted_class_name']} ({result['confidence']:.3f})")
-                print("✅ Handler fonctionnel!")
+                print("✅ Handler V3 ULTRA LIGHT fonctionnel!")
             else:
                 print(f"❌ Erreur: {result.get('error', 'Inconnue')}")
         else:
-            print("❌ Échec de l'initialisation")
+            print("❌ Échec initialisation")
             
     except Exception as e:
-        print(f"❌ Erreur de test: {e}")
-        print(f"🔍 Traceback: {traceback.format_exc()}") 
+        print(f"❌ Erreur test: {e}")