Update handler.py

handler.py

@@ -5,11 +5,12 @@ import base64
 import io
 import torch.nn.functional as F
 import gc
+import json
 
 class EndpointHandler:
     def __init__(self, path=""):
-        print("🚀 VerifAI Handler V4 MINIMAL - Initialisation")
-        print("⚡ Version ultra-simplifiée")
+        print("🚀 VerifAI Handler V5 ULTRA ROBUST - Initialisation")
+        print("⚡ Version ultra-robuste")
         
         self.model = None
         self.processor = None
@@ -19,7 +20,6 @@ class EndpointHandler:
         try:
             print("🔄 Chargement modèle...")
             
-            # Chargement simple
             self.processor = AutoImageProcessor.from_pretrained(self.model_name)
             self.model = AutoModelForImageClassification.from_pretrained(
                 self.model_name,
@@ -27,7 +27,6 @@ class EndpointHandler:
             )
             self.model.eval()
             
-            # Labels
             if hasattr(self.model.config, 'id2label'):
                 self.model_labels = self.model.config.id2label
             else:
@@ -35,14 +34,54 @@ class EndpointHandler:
             
             print("✅ Modèle chargé")
             print(f"📋 Labels: {self.model_labels}")
-            print("🎯 Handler V4 prêt!")
+            print("🎯 Handler V5 prêt!")
             
         except Exception as e:
             print(f"❌ Erreur: {e}")
             self.model = None
             self.processor = None
     
-    def _normalize_label(self, label: str) -> str:
+    def _extract_image_data(self, data):
+        """Extraction robuste des données image"""
+        try:
+            # Cas 1: data est directement la string base64
+            if isinstance(data, str):
+                print("📄 Input détecté: string directe")
+                return data
+            
+            # Cas 2: data est un dict avec clé "inputs"
+            if isinstance(data, dict):
+                print("📄 Input détecté: dictionnaire")
+                
+                # Essayer "inputs"
+                if "inputs" in data:
+                    return data["inputs"]
+                
+                # Essayer d'autres clés communes
+                for key in ["image", "data", "input", "content"]:
+                    if key in data:
+                        return data[key]
+                
+                # Si aucune clé connue, prendre la première valeur
+                if data:
+                    first_value = list(data.values())[0]
+                    print(f"🔍 Utilisation de la première valeur: {type(first_value)}")
+                    return first_value
+            
+            # Cas 3: data est une liste
+            if isinstance(data, list) and len(data) > 0:
+                print("📄 Input détecté: liste")
+                return data[0]
+            
+            # Cas 4: autres types
+            print(f"📄 Input détecté: {type(data)}")
+            return str(data)
+            
+        except Exception as e:
+            print(f"⚠️  Erreur extraction: {e}")
+            return None
+    
+    def _normalize_label(self, label):
         """Normalise les labels"""
         if not isinstance(label, str):
             label = str(label)
@@ -65,24 +104,38 @@ class EndpointHandler:
                 "confidence": 0.0,
                 "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
                 "cam_image": None,
-                "version": "4.0-minimal",
-                "handler_name": "VerifAI Handler V4 MINIMAL"
+                "version": "5.0-ultra-robust",
+                "handler_name": "VerifAI Handler V5 ULTRA ROBUST"
             }
         
         try:
-            print("🔄 Traitement...")
+            print("🔄 Traitement ultra-robuste...")
+            print(f"🔍 Type d'input reçu: {type(data)}")
+            
+            # Extraction robuste des données
+            image_data = self._extract_image_data(data)
             
-            # Extraction image
-            image_data = data.get("inputs") or data
             if not image_data:
-                raise ValueError("Pas de données image")
+                raise ValueError("Aucune donnée image trouvée")
             
-            # Décodage
-            if isinstance(image_data, str) and image_data.startswith('data:'):
-                image_data = image_data.split(',', 1)[1]
+            print(f"🔍 Données extraites: {type(image_data)}, longueur: {len(str(image_data)) if image_data else 0}")
             
-            image_bytes = base64.b64decode(image_data)
-            image = Image.open(io.BytesIO(image_bytes))
+            # Nettoyage du base64
+            if isinstance(image_data, str):
+                # Supprimer le préfixe data URI si présent
+                if image_data.startswith('data:'):
+                    image_data = image_data.split(',', 1)[1]
+                
+                # Supprimer les espaces et retours de ligne
+                image_data = image_data.strip().replace('\n', '').replace('\r', '').replace(' ', '')
+            
+            # Décodage
+            try:
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(io.BytesIO(image_bytes))
+                print(f"✅ Image décodée: {image.size}, mode: {image.mode}")
+            except Exception as e:
+                raise ValueError(f"Erreur décodage base64: {e}")
             
             if image.mode != 'RGB':
                 image = image.convert('RGB')
@@ -90,6 +143,7 @@ class EndpointHandler:
             # Redimensionnement si nécessaire
             if image.size[0] * image.size[1] > 1048576:  # Plus de 1MP
                 image = image.resize((512, 512), Image.Resampling.LANCZOS)
+                print("⚠️  Image redimensionnée")
 
             print("🧠 Inférence...")
             
@@ -131,17 +185,23 @@ class EndpointHandler:
                 "predicted_class_name": prediction_label,
                 "confidence": confidence,
                 "class_probabilities": class_probs,
-                "cam_image": None,  # Désactivé
+                "cam_image": None,
                 "model_info": {
                     "model_name": self.model_name,
-                    "handler_version": "verifai-v4-minimal",
+                    "handler_version": "verifai-v5-ultra-robust",
                     "precision_mode": "fast",
                     "raw_prediction_id": predicted_class_id,
                     "raw_labels": self.model_labels
                 },
-                "version": "4.0-minimal",
-                "handler_name": "VerifAI Handler V4 MINIMAL",
-                "note": "Version ultra-simplifiée pour compatibilité maximale"
+                "version": "5.0-ultra-robust",
+                "handler_name": "VerifAI Handler V5 ULTRA ROBUST",
+                "note": "Version ultra-robuste - gère tous les formats d'entrée",
+                "input_analysis": {
+                    "original_type": str(type(data)),
+                    "extracted_type": str(type(image_data)),
+                    "image_size": image.size,
+                    "image_mode": image.mode
+                }
             }
 
         except Exception as e:
@@ -155,14 +215,18 @@ class EndpointHandler:
                 "confidence": 0.0,
                 "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
                 "cam_image": None,
-                "version": "4.0-minimal",
-                "handler_name": "VerifAI Handler V4 MINIMAL"
+                "version": "5.0-ultra-robust",
+                "handler_name": "VerifAI Handler V5 ULTRA ROBUST",
+                "debug_info": {
+                    "input_type": str(type(data)),
+                    "input_content": str(data)[:100] + "..." if data else "None"
+                }
             }
 
-# Test simple
+# Test
 if __name__ == "__main__":
-    print("🧪 TEST HANDLER V4 MINIMAL")
-    print("=" * 40)
+    print("🧪 TEST HANDLER V5 ULTRA ROBUST")
+    print("=" * 50)
     
     try:
         handler = EndpointHandler()
@@ -170,20 +234,29 @@ if __name__ == "__main__":
         if handler.model is not None:
             print("✅ Initialisation OK")
             
-            # Test
+            # Test avec différents formats
             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})
+            test_cases = [
+                {"inputs": test_data},  # Format dict
+                test_data,  # String directe
+                [test_data],  # Liste
+            ]
             
-            print(f"📊 Statut: {result['status']}")
-            if result['status'] == 'success':
-                print(f"🎯 Prédiction: {result['predicted_class_name']} ({result['confidence']:.3f})")
-                print("✅ Handler V4 MINIMAL fonctionnel!")
-            else:
-                print(f"❌ Erreur: {result.get('error', 'Inconnue')}")
+            for i, test_case in enumerate(test_cases, 1):
+                print(f"\n🔄 Test {i}: {type(test_case)}")
+                result = handler(test_case)
+                
+                print(f"📊 Statut: {result['status']}")
+                if result['status'] == 'success':
+                    print(f"🎯 Prédiction: {result['predicted_class_name']} ({result['confidence']:.3f})")
+                else:
+                    print(f"❌ Erreur: {result.get('error', 'Inconnue')}")
+            
+            print("\n✅ Handler V5 ULTRA ROBUST testé!")
         else:
             print("❌ Échec initialisation")