Update handler.py

handler.py

@@ -1,50 +1,95 @@
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from transformers import ViTFeatureExtractor, ViTForImageClassification
 from PIL import Image
 import torch
+import base64
+import io
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        self.model = AutoModelForImageClassification.from_pretrained(model_dir)
-        self.processor = AutoFeatureExtractor.from_pretrained(model_dir)
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
+        """
+        تهيئة النموذج ومعالج الميزات
         
+        المعلمات:
+            model_dir: مسار مجلد النموذج
+        """
+        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
+    
     def __call__(self, data):
         """
-        Args:
-            data: Image data in binary format
-        Returns:
-            Prediction result as a dictionary
+        معالجة البيانات المدخلة وإجراء التنبؤ
+        
+        المعلمات:
+            data: بيانات الصورة المشفرة بـ base64 أو كائن الصورة
+            
+        العائد:
+            dict: نتائج التنبؤ
         """
         try:
-            # Load and process image
-            image = Image.open(data).convert("RGB")
-            inputs = self.processor(images=image, return_tensors="pt")
+            # التحقق من نوع البيانات المدخلة
+            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]
+                
+                # فك تشفير البيانات
+                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")
+            else:
+                return {"error": "تنسيق البيانات غير مدعوم"}
+            
+            # معالجة الصورة
+            inputs = self.feature_extractor(images=image, return_tensors="pt")
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
-            # Get prediction
+            # إجراء التنبؤ
             with torch.no_grad():
                 outputs = self.model(**inputs)
                 logits = outputs.logits
                 probabilities = torch.nn.functional.softmax(logits, dim=-1)
                 
-            # Get prediction class and confidence
-            predicted_class_idx = probabilities.argmax().item()
-            confidence = probabilities[0][predicted_class_idx].item()
-            
-            # Get class labels
-            id2label = self.model.config.id2label
-            predicted_class = id2label[predicted_class_idx]
-            
-            # Return results
-            return {
-                "predicted_class": predicted_class,
-                "confidence": confidence,
-                "all_probabilities": {
-                    id2label[i]: prob.item() 
-                    for i, prob in enumerate(probabilities[0])
+                # الحصول على التصنيف ونسبة الثقة
+                predicted_class_idx = probabilities.argmax().item()
+                confidence = probabilities[0][predicted_class_idx].item()
+                
+                # تحويل الفهرس إلى تسمية
+                id2label = self.model.config.id2label
+                predicted_class = id2label[predicted_class_idx]
+            
+            # إعداد النتائج
+            results = {
+                "prediction": predicted_class,
+                "confidence": float(confidence),
+                "is_fake": predicted_class == "fake",
+                "probabilities": {
+                    label: float(prob)
+                    for label, prob in zip(id2label.values(), probabilities[0].cpu().numpy())
                 }
             }
             
+            return results
+        
         except Exception as e:
-            return {"error": str(e)}
+            # معالجة الأخطاء
+            error_message = str(e)
+            print(f"خطأ في معالجة الصورة: {error_message}")
+            return {"error": error_message}