Create handler.py

handler.py

@@ -0,0 +1,49 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import os
+import json
+
+# لو النموذج PyTorch .pth
+class AIImageSourceHandler:
+    def __init__(self):
+        self.model = None
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.labels = ['Real', 'Midjourney', 'DALL·E', 'StableDiffusion']  # عدّل حسب عدد التصنيفات
+
+    def initialize(self, model_dir: str):
+        # تحميل النموذج
+        model_path = os.path.join(model_dir, "model.pth")  # عدّل الاسم لو مختلف
+        self.model = torch.load(model_path, map_location=self.device)
+        self.model.eval()
+
+        # إعداد التحويلات للصورة
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+        ])
+
+    def preprocess(self, image: Image.Image):
+        image = image.convert("RGB")
+        return self.transform(image).unsqueeze(0).to(self.device)
+
+    def predict(self, inputs):
+        image = inputs.get("image")
+        if image is None:
+            return {"error": "No image provided"}
+
+        img_tensor = self.preprocess(image)
+        with torch.no_grad():
+            outputs = self.model(img_tensor)
+            probs = torch.nn.functional.softmax(outputs[0], dim=0)
+
+        results = {
+            label: float(probs[i]) for i, label in enumerate(self.labels)
+        }
+        return results
+
+    def __call__(self, data):
+        image_data = data.get("inputs") or data.get("image")
+        if isinstance(image_data, Image.Image):
+            return self.predict({"image": image_data})
+        return {"error": "Invalid input"}