File size: 5,972 Bytes
dd32056 54240fb dd32056 00fa5d2 963149c 00fa5d2 8ea56b1 dd32056 bc35f37 00fa5d2 bc35f37 8b1e242 963149c bc35f37 963149c bc35f37 963149c bc35f37 8b1e242 dd32056 00fa5d2 ba3f2d0 bc35f37 ba3f2d0 54240fb bc35f37 dd32056 ba3f2d0 bc35f37 963149c bc35f37 963149c bc35f37 54240fb bc35f37 9a050d4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import base64
import io
import os
from typing import Dict, Any, List
import torch
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification
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}")
# استخدام مكتبة transformers بدلاً من timm
try:
print("تحميل معالج الصور ViT")
self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
print("تحميل نموذج ViT")
self.model = ViTForImageClassification.from_pretrained(
"google/vit-base-patch16-224",
num_labels=5,
id2label={
0: "stable_diffusion",
1: "midjourney",
2: "dalle",
3: "real",
4: "other_ai"
},
label2id={
"stable_diffusion": 0,
"midjourney": 1,
"dalle": 2,
"real": 3,
"other_ai": 4
}
)
# محاولة تحميل الأوزان المخصصة إذا كانت موجودة
pytorch_path = os.path.join(model_dir, "pytorch_model.bin")
if os.path.exists(pytorch_path):
print(f"محاولة تحميل الأوزان المخصصة من: {pytorch_path}")
try:
state_dict = torch.load(pytorch_path, map_location="cpu")
# تحميل الأوزان المتوافقة فقط
self.model.load_state_dict(state_dict, strict=False)
print("تم تحميل الأوزان المخصصة بنجاح")
except Exception as e:
print(f"تحذير: فشل تحميل الأوزان المخصصة: {e}")
self.model.to(self.device)
self.model.eval()
print("تم تهيئة النموذج بنجاح")
except Exception as e:
print(f"خطأ في تهيئة النموذج: {e}")
import traceback
traceback.print_exc()
raise
self.labels = [
"stable_diffusion",
"midjourney",
"dalle",
"real",
"other_ai",
]
def _decode_b64(self, b: bytes) -> Image.Image:
try:
print(f"فك ترميز base64. حجم البيانات: {len(b)} بايت")
img = Image.open(io.BytesIO(base64.b64decode(b)))
print(f"تم فك الترميز بنجاح. حجم الصورة: {img.size}, وضع الصورة: {img.mode}")
return img
except Exception as e:
print(f"خطأ في فك ترميز base64: {e}")
raise
def __call__(self, data: Any) -> List[Dict[str, Any]]:
print(f"استدعاء __call__ مع البيانات من النوع: {type(data)}")
img: Image.Image | None = None
try:
if isinstance(data, Image.Image):
print("البيانات هي صورة PIL")
img = data
elif isinstance(data, dict):
print(f"البيانات هي قاموس بالمفاتيح: {list(data.keys())}")
payload = data.get("inputs") or data.get("image")
print(f"نوع الحمولة: {type(payload)}")
if isinstance(payload, (str, bytes)):
if isinstance(payload, str):
print("تحويل السلسلة النصية إلى بايت")
payload = payload.encode()
img = self._decode_b64(payload)
if img is None:
print("لم يتم العثور على صورة صالحة في البيانات")
return [{"label": "error", "score": 1.0}]
print("معالجة الصورة باستخدام معالج ViT")
inputs = self.processor(images=img, return_tensors="pt").to(self.device)
print("بدء التنبؤ باستخدام النموذج")
with torch.no_grad():
outputs = self.model(**inputs)
logits = outputs.logits
probs = torch.nn.functional.softmax(logits[0], dim=0)
print(f"تم الحصول على الاحتمالات: {probs}")
# تحويل النتائج إلى التنسيق المطلوب: Array<label: string, score:number>
results = []
for i, label in enumerate(self.labels):
score = float(probs[i])
print(f"التسمية: {label}, الدرجة: {score}")
results.append({
"label": label,
"score": score
})
# ترتيب النتائج تنازلياً حسب درجة الثقة
results.sort(key=lambda x: x["score"], reverse=True)
print(f"النتائج النهائية: {results}")
return results
except Exception as e:
print(f"حدث خطأ أثناء المعالجة: {e}")
print(f"نوع الخطأ: {type(e).__name__}")
print(f"تفاصيل الخطأ: {str(e)}")
import traceback
traceback.print_exc()
return [{"label": "error", "score": 1.0}]
|