Update handler.py

handler.py

@@ -1,95 +1,108 @@
-from transformers import ViTFeatureExtractor, ViTForImageClassification
-from PIL import Image
+import os
 import torch
-import base64
-import io
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+import json
 
 class EndpointHandler:
     def __init__(self, model_dir):
         """
-        تهيئة النموذج ومعالج الميزات
-        
-        المعلمات:
-            model_dir: مسار مجلد النموذج
+        Initialize the model for AI image detection
         """
-        try:
-            # تحميل النموذج ومعالج الميزات
-            self.feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir)
-            self.model = ViTForImageClassification.from_pretrained(model_dir)
-            
-            # نقل النموذج إلى وحدة المعالجة المناسبة
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            self.model.to(self.device)
-            
-            # وضع النموذج في وضع التقييم
-            self.model.eval()
-            
-            print(f"تم تحميل النموذج بنجاح على جهاز {self.device}")
-        except Exception as e:
-            print(f"خطأ في تحميل النموذج: {str(e)}")
-            raise
+        # Set device
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Load model
+        self.model = self._load_model(model_dir)
+        
+        # Define transforms
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+        
+        # Class names
+        self.classes = ["Real Image", "AI-Generated Image"]
+    
+    def _load_model(self, model_dir):
+        # Create model architecture
+        model = models.efficientnet_v2_s(weights=None)
+        
+        # Recreate classifier exactly as in training
+        model.classifier = nn.Sequential(
+            nn.Linear(model.classifier[1].in_features, 1024),
+            nn.ReLU(),
+            nn.Dropout(p=0.3),
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Dropout(p=0.3),
+            nn.Linear(512, 2)
+        )
+        
+        # Load state dict - find the pth file in the directory
+        model_path = os.path.join(model_dir, "best_model_improved.pth")
+        
+        if os.path.exists(model_path):
+            print(f"Loading model from {model_path}")
+            model.load_state_dict(torch.load(model_path, map_location=self.device))
+            model.to(self.device)
+            model.eval()
+            return model
+        else:
+            raise FileNotFoundError(f"Model file not found at {model_path}. Files in directory: {os.listdir(model_dir)}")
     
     def __call__(self, data):
         """
-        معالجة البيانات المدخلة وإجراء التنبؤ
-        
-        المعلمات:
-            data: بيانات الصورة المشفرة بـ base64 أو كائن الصورة
-            
-        العائد:
-            dict: نتائج التنبؤ
+        Run prediction on the input data
         """
         try:
-            # التحقق من نوع البيانات المدخلة
-            if isinstance(data, dict) and "image" in data:
-                # إذا كانت البيانات مشفرة بـ base64
-                image_data = data["image"]
-                if isinstance(image_data, str) and image_data.startswith("data:image"):
-                    # إزالة بادئة URL للبيانات
-                    image_data = image_data.split(",")[1]
+            # Parse request data
+            if isinstance(data, dict) and "inputs" in data:
+                # API format
+                input_data = data["inputs"]
+            else:
+                # Direct image
+                input_data = data
+            
+            # Process image
+            if isinstance(input_data, str):  # Base64 string
+                import base64
+                from io import BytesIO
                 
-                # فك تشفير البيانات
-                image_bytes = base64.b64decode(image_data)
-                image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-            elif isinstance(data, bytes):
-                # إذا كانت البيانات ثنائية
-                image = Image.open(io.BytesIO(data)).convert("RGB")
+                # Decode base64 image
+                image_bytes = base64.b64decode(input_data.split(",")[-1])
+                image = Image.open(BytesIO(image_bytes)).convert("RGB")
+            elif hasattr(input_data, "read"):  # File-like object
+                image = Image.open(input_data).convert("RGB")
             else:
-                return {"error": "تنسيق البيانات غير مدعوم"}
+                # Assume PIL Image
+                image = input_data
             
-            # معالجة الصورة
-            inputs = self.feature_extractor(images=image, return_tensors="pt")
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            # Preprocess image
+            image_tensor = self.transform(image).unsqueeze(0).to(self.device)
             
-            # إجراء التنبؤ
+            # Make prediction
             with torch.no_grad():
-                outputs = self.model(**inputs)
-                logits = outputs.logits
-                probabilities = torch.nn.functional.softmax(logits, dim=-1)
-                
-                # الحصول على التصنيف ونسبة الثقة
-                predicted_class_idx = probabilities.argmax().item()
-                confidence = probabilities[0][predicted_class_idx].item()
-                
-                # تحويل الفهرس إلى تسمية
-                id2label = self.model.config.id2label
-                predicted_class = id2label[predicted_class_idx]
+                outputs = self.model(image_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
+                prediction = torch.argmax(probabilities).item()
             
-            # إعداد النتائج
-            results = {
-                "prediction": predicted_class,
-                "confidence": float(confidence),
-                "is_fake": predicted_class == "fake",
+            # Format results
+            real_prob = probabilities[0].item() * 100
+            ai_prob = probabilities[1].item() * 100
+            
+            result = {
+                "prediction": self.classes[prediction],
+                "confidence": float(probabilities[prediction].item()),
                 "probabilities": {
-                    label: float(prob)
-                    for label, prob in zip(id2label.values(), probabilities[0].cpu().numpy())
+                    "Real Image": float(real_prob),
+                    "AI-Generated Image": float(ai_prob)
                 }
             }
             
-            return results
-        
+            return result
+            
         except Exception as e:
-            # معالجة الأخطاء
-            error_message = str(e)
-            print(f"خطأ في معالجة الصورة: {error_message}")
-            return {"error": error_message}
+            return {"error": str(e)}