vaibhav07112004's picture
Update api/models/face_models.py
66e212d verified
import torch
from app.models.model_definitions import create_aadcn_model
from app.config import FACE_MODEL_PATH
from deepface import DeepFace
from PIL import Image
import io
import numpy as np
from torchvision import transforms
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
face_model = create_aadcn_model(num_classes=8)
face_model.load_state_dict(torch.load(FACE_MODEL_PATH, map_location=device))
face_model.eval()
emotion_idx_to_label = {
0: 'angry',
1: 'contempt',
2: 'disgust',
3: 'fear',
4: 'happy',
5: 'neutral',
6: 'sad',
7: 'surprise'
}
dataset_to_custom = {
'angry': ['Anger', 'Annoyance', 'Disapproval'],
'contempt': ['Disapproval', 'Disconnection', 'Annoyance'],
'disgust': ['Aversion', 'Disapproval', 'Disconnection'],
'fear': ['Fear', 'Disquietment', 'Doubt/Confusion'],
'happy': ['Happiness', 'Affection', 'Pleasure', 'Excitement'],
'neutral': ['Peace', 'Esteem', 'Confidence'],
'sad': ['Sadness', 'Fatigue', 'Suffering'],
'surprise': ['Surprise', 'Anticipation', 'Excitement']
}
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])
def predict_face_emotion_hybrid(image_bytes):
# Try DeepFace for "happy" detection first
try:
image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
np_image = np.array(image) # Ensure numpy array is passed
result = DeepFace.analyze(
np_image,
actions=['emotion'],
enforce_detection=False,
detector_backend='opencv'
)
if isinstance(result, list):
result = result[0]
if result['dominant_emotion'] == 'happy':
return 'happy', dataset_to_custom['happy']
except Exception as e:
print(f"DeepFace error: {e}. Falling back to AA-DCN model.")
# Fallback to AA-DCN
image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
img_t = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
output = face_model(img_t)
pred_idx = output.argmax(dim=1).item()
dataset_label = emotion_idx_to_label.get(pred_idx, 'unknown')
custom_moods = dataset_to_custom.get(dataset_label, ['Unknown'])
return dataset_label, custom_moods