File size: 1,634 Bytes
eb9ae8c | 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 | import torch
import base64
import os
import io
from PIL import Image
from transformers import AutoImageProcessor, AutoConfig, AutoModelForImageClassification
class EmotionEngine:
def __init__(self):
# اسم الموديل مش مهم يظهر في أي حتة تانية
self.processor = AutoImageProcessor.from_pretrained(
"trpakov/vit-face-expression"
)
config = AutoConfig.from_pretrained(
"trpakov/vit-face-expression"
)
self.model = AutoModelForImageClassification.from_config(config)
# ... داخل الكلاس __init__
model_dir = os.path.join(os.path.dirname(__file__), "trained_models")
model_path = os.path.join(model_dir, "emotion_model.pth")
state_dict = torch.load(model_path, map_location="cpu")
self.model.load_state_dict(state_dict)
self.model.eval()
self.labels = self.model.config.id2label
def predict_from_base64(self, base64_img):
# فك الصورة
img_bytes = base64.b64decode(base64_img.split(",")[1])
img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
# preprocessing
inputs = self.processor(images=img, return_tensors="pt")
with torch.no_grad():
outputs = self.model(**inputs)
probs = torch.softmax(outputs.logits, dim=1)[0]
emotion_probs = {
self.labels[i]: float(probs[i])
for i in range(len(probs))
}
dominant_emotion = max(
emotion_probs, key=emotion_probs.get
)
return emotion_probs, dominant_emotion
|