yhay360 commited on
Commit
c1dc91f
·
1 Parent(s): 73493c1

chore: remove custom handler – use default pipeline

Browse files
Files changed (1) hide show
  1. handler.py +0 -49
handler.py DELETED
@@ -1,49 +0,0 @@
1
- import torch
2
- from torchvision import transforms
3
- from PIL import Image
4
- import os
5
- import json
6
-
7
- # لو النموذج PyTorch .pth
8
- class AIImageSourceHandler:
9
- def __init__(self):
10
- self.model = None
11
- self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
- self.labels = ['Real', 'Midjourney', 'DALL·E', 'StableDiffusion'] # عدّل حسب عدد التصنيفات
13
-
14
- def initialize(self, model_dir: str):
15
- # تحميل النموذج
16
- model_path = os.path.join(model_dir, "model.pth") # عدّل الاسم لو مختلف
17
- self.model = torch.load(model_path, map_location=self.device)
18
- self.model.eval()
19
-
20
- # إعداد التحويلات للصورة
21
- self.transform = transforms.Compose([
22
- transforms.Resize((224, 224)),
23
- transforms.ToTensor(),
24
- ])
25
-
26
- def preprocess(self, image: Image.Image):
27
- image = image.convert("RGB")
28
- return self.transform(image).unsqueeze(0).to(self.device)
29
-
30
- def predict(self, inputs):
31
- image = inputs.get("image")
32
- if image is None:
33
- return {"error": "No image provided"}
34
-
35
- img_tensor = self.preprocess(image)
36
- with torch.no_grad():
37
- outputs = self.model(img_tensor)
38
- probs = torch.nn.functional.softmax(outputs[0], dim=0)
39
-
40
- results = {
41
- label: float(probs[i]) for i, label in enumerate(self.labels)
42
- }
43
- return results
44
-
45
- def __call__(self, data):
46
- image_data = data.get("inputs") or data.get("image")
47
- if isinstance(image_data, Image.Image):
48
- return self.predict({"image": image_data})
49
- return {"error": "Invalid input"}