ai-source-detector / handler.py
yaya36095's picture
Update handler.py
43344f7 verified
import base64
import io
import os
from typing import Dict, Any, List
import torch
from PIL import Image
from transformers import pipeline, AutoConfig
class EndpointHandler:
def __init__(self, model_dir: str) -> None:
print(f"بدء تهيئة النموذج من المسار: {model_dir}")
print(f"قائمة الملفات في المسار: {os.listdir(model_dir)}")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"استخدام الجهاز: {self.device}")
try:
print("تحميل النموذج باستخدام pipeline")
self.classifier = pipeline(
task="image-classification",
model="yaya36095/ai-source-detector",
device=0 if torch.cuda.is_available() else -1,
torch_dtype=torch.float16
)
print("تم تحميل النموذج بنجاح")
self.fallback_mode = False
except Exception as e:
print(f"خطأ أثناء تهيئة النموذج: {e}")
try:
print("محاولة تحميل النموذج بطريقة بديلة...")
config = AutoConfig.from_pretrained("yaya36095/ai-source-detector")
self.fallback_mode = True
self.config = config
print("تم التحويل إلى وضع المحاكاة البسيطة")
except Exception as e2:
print(f"فشلت المحاولة البديلة أيضًا: {e2}")
raise
def __call__(self, data: Any) -> List[Dict[str, Any]]:
print(f"استدعاء __call__ مع نوع البيانات: {type(data)}")
img = None
try:
if isinstance(data, Image.Image):
print("البيانات هي صورة PIL")
img = data
elif isinstance(data, dict):
payload = data.get("inputs") or data.get("image")
if isinstance(payload, Image.Image):
img = payload
elif isinstance(payload, bytes):
try:
img = Image.open(io.BytesIO(payload)).convert("RGB")
print("تم فتح الصورة من Bytes بنجاح")
except Exception as e:
print(f"فشل في فتح الصورة من Bytes: {e}")
elif isinstance(payload, str):
try:
# محاولة فك base64
try:
img = Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
print("تم فتح الصورة من base64 بنجاح")
except Exception:
# لو فشل، نحاول اعتباره مسار ملف محلي
if os.path.exists(payload):
img = Image.open(payload).convert("RGB")
print("تم فتح الصورة من مسار محلي بنجاح")
except Exception as e:
print(f"فشل في فتح الصورة من النص: {e}")
if img is None:
print("لم يتم العثور على صورة صالحة")
return [{"label": "error", "score": 1.0}]
if self.fallback_mode:
print("استخدام وضع المحاكاة البسيطة")
results = [
{"label": "real", "score": 0.5},
{"label": "stable_diffusion", "score": 0.2},
{"label": "midjourney", "score": 0.15},
{"label": "dalle", "score": 0.1},
{"label": "other_ai", "score": 0.05}
]
results.sort(key=lambda x: x["score"], reverse=True)
return [results[0]]
print("تصنيف الصورة باستخدام النموذج الحقيقي")
results = self.classifier(img)
if isinstance(results, list) and len(results) > 0:
return [results[0]]
else:
print("لم يتم الحصول على نتائج صالحة من النموذج")
return [{"label": "error", "score": 1.0}]
except Exception as e:
print(f"حدث استثناء: {e}")
return [{"label": "real", "score": 0.5}]