|
|
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 _decode_b64(self, b: bytes) -> Image.Image: |
|
|
try: |
|
|
print(f"فك ترميز base64. حجم البيانات: {len(b)} بايت") |
|
|
return Image.open(io.BytesIO(base64.b64decode(b))).convert("RGB") |
|
|
except Exception as e: |
|
|
print(f"خطأ في فك الترميز: {e}") |
|
|
raise |
|
|
|
|
|
def __call__(self, data: Any) -> List[Dict[str, Any]]: |
|
|
print(f"استدعاء __call__ مع نوع البيانات: {type(data)}") |
|
|
|
|
|
img = None |
|
|
try: |
|
|
if isinstance(data, Image.Image): |
|
|
img = data |
|
|
elif isinstance(data, dict): |
|
|
payload = data.get("inputs") or data.get("image") |
|
|
if isinstance(payload, str): |
|
|
payload = payload.encode() |
|
|
if isinstance(payload, bytes): |
|
|
img = self._decode_b64(payload) |
|
|
|
|
|
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}] |
|
|
|
