| 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) |
|
|
| |
| 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") |
|
|
| |
| 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 |
|
|
