handler.py

import torch
from transformers import AutoModelForImageClassification, AutoImageProcessor
from PIL import Image
import base64
import io
import torch.nn.functional as F
import gc
import json

class EndpointHandler:
    def __init__(self, path=""):
        print("🚀 VerifAI Handler V5 ULTRA ROBUST - Initialisation")
        print("⚡ Version ultra-robuste")
        
        self.model = None
        self.processor = None
        self.model_labels = {}
        self.model_name = "haywoodsloan/ai-image-detector-deploy"
        
        try:
            print("🔄 Chargement modèle...")
            
            self.processor = AutoImageProcessor.from_pretrained(self.model_name)
            self.model = AutoModelForImageClassification.from_pretrained(
                self.model_name,
                torch_dtype=torch.float32
            )
            self.model.eval()
            
            if hasattr(self.model.config, 'id2label'):
                self.model_labels = self.model.config.id2label
            else:
                self.model_labels = {0: "Real", 1: "Fake"}
            
            print("✅ Modèle chargé")
            print(f"📋 Labels: {self.model_labels}")
            print("🎯 Handler V5 prêt!")
            
        except Exception as e:
            print(f"❌ Erreur: {e}")
            self.model = None
            self.processor = None
    
    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)
        
        label_lower = label.lower()
        if any(word in label_lower for word in ['real', 'human', 'authentic']):
            return "Human"
        if any(word in label_lower for word in ['fake', 'generated', 'ai', 'artificial']):
            return "AI Generated"
        return "Unknown"
    
    def __call__(self, data):
        # Vérification
        if self.model is None or self.processor is None:
            return {
                "status": "error",
                "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": "5.0-ultra-robust",
                "handler_name": "VerifAI Handler V5 ULTRA ROBUST"
            }
        
        try:
            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)
            
            if not image_data:
                raise ValueError("Aucune donnée image trouvée")
            
            print(f"🔍 Données extraites: {type(image_data)}, longueur: {len(str(image_data)) if image_data else 0}")
            
            # 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')
            
            # 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...")
            
            # Inférence
            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()
            
            # Résultats
            class_probs = {}
            for class_id, prob in enumerate(probabilities):
                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)

            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)
            prediction_id = 1 if prediction_label == "AI Generated" else 0
            
            print(f"🎯 Résultat: {prediction_label} ({confidence:.3f})")

            # Nettoyage
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            gc.collect()

            return {
                "status": "success",
                "prediction": prediction_id,
                "predicted_class_name": prediction_label,
                "confidence": confidence,
                "class_probabilities": class_probs,
                "cam_image": None,
                "model_info": {
                    "model_name": self.model_name,
                    "handler_version": "verifai-v5-ultra-robust",
                    "precision_mode": "fast",
                    "raw_prediction_id": predicted_class_id,
                    "raw_labels": self.model_labels
                },
                "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:
            print(f"❌ Erreur: {e}")
            
            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": "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
if __name__ == "__main__":
    print("🧪 TEST HANDLER V5 ULTRA ROBUST")
    print("=" * 50)
    
    try:
        handler = EndpointHandler()
        
        if handler.model is not None:
            print("✅ Initialisation OK")
            
            # 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')
            
            test_cases = [
                {"inputs": test_data},  # Format dict
                test_data,  # String directe
                [test_data],  # Liste
            ]
            
            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")
            
    except Exception as e:
        print(f"❌ Erreur test: {e}")