Update handler.py

handler.py

@@ -1,37 +1,25 @@
-import torch
-from torchvision import transforms
-from PIL import Image
-from torchvision.models import efficientnet_v2_s
-
-class EndpointHandler:
-    def __init__(self, path=""):
-        self.model = efficientnet_v2_s(weights=None)
-        self.model.classifier = torch.nn.Sequential(
-            torch.nn.Linear(self.model.classifier[1].in_features, 1024),
-            torch.nn.ReLU(),
-            torch.nn.Dropout(0.3),
-            torch.nn.Linear(1024, 512),
-            torch.nn.ReLU(),
-            torch.nn.Dropout(0.3),
-            torch.nn.Linear(512, 2)
-        )
-        self.model.load_state_dict(torch.load(f"{path}/pytorch_model.bin", map_location=torch.device("cpu")))
-        self.model.eval()
-
-        self.transform = transforms.Compose([
-            transforms.Resize((224, 224)),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406],
-                                 [0.229, 0.224, 0.225])
-        ])
-
-    def __call__(self, data):
-        image = Image.open(data["inputs"]).convert("RGB")
-        image = self.transform(image).unsqueeze(0)
-        with torch.no_grad():
-            outputs = self.model(image)
-            probs = torch.nn.functional.softmax(outputs[0], dim=0)
-            return {
-                "real": round(probs[0].item(), 4),
-                "fake": round(probs[1].item(), 4)
-            }
+from typing import Dict, Any
+from PIL import Image
+from io import BytesIO
+import base64
+
+from inference import load_model, predict
+
+# تحميل النموذج عند بدء السيرفر
+model = load_model()
+
+# الدالة اللي بيتم استدعاؤها عند رفع صورة
+def predict_image(inputs: Dict[str, Any]) -> Dict[str, float]:
+    # لو كانت الصورة مرسلة كـ base64
+    if "image" in inputs:
+        image_data = inputs["image"]
+        if isinstance(image_data, str):
+            image = Image.open(BytesIO(base64.b64decode(image_data))).convert("RGB")
+        else:
+            image = Image.open(image_data).convert("RGB")
+    else:
+        raise ValueError("Missing 'image' key in input")
+
+    # التنبؤ
+    result = predict(model, image)
+    return result