DrGM-FastViT-nano-Multimodal-Facial-Emotion-Recognition

import torch
import timm
import mediapipe as mp
import numpy as np
from PIL import Image
from torchvision import transforms

# 1. Define Model Architecture (Same as training)
class MultimodalFERModel(torch.nn.Module):
    def __init__(self, num_classes=7):
        super().__init__()
        self.vision_backbone = timm.create_model('fastvit_sa24.apple_in1k', num_classes=0)
        self.landmark_encoder = torch.nn.Sequential(
            torch.nn.Linear(478*3, 512), torch.nn.BatchNorm1d(512), torch.nn.ReLU(),
            torch.nn.Linear(512, 256), torch.nn.BatchNorm1d(256), torch.nn.ReLU()
        )
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(self.vision_backbone.num_features + 256, 512),
            torch.nn.BatchNorm1d(512), torch.nn.ReLU(),
            torch.nn.Linear(512, num_classes)
        )
    def forward(self, pixel_values, landmarks):
        v = self.vision_backbone(pixel_values)
        l = self.landmark_encoder(landmarks)
        return self.classifier(torch.cat((v, l), dim=1))

# 2. Load Model
model = MultimodalFERModel()
# Download weights from Hub (manually or via API) and load
# model.load_state_dict(torch.hub.load_state_dict_from_url('...'))
model.eval()

# 3. Preprocessing (MediaPipe + Transforms)
mp_face = mp.solutions.face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1)
val_tf = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def predict(image_path):
    img = Image.open(image_path).convert('RGB')
    # Vision Input
    pixel_values = val_tf(img).unsqueeze(0)
    
    # Landmark Input
    results = mp_face.process(np.array(img))
    if results.multi_face_landmarks:
        lm = np.array([[l.x, l.y, l.z] for l in results.multi_face_landmarks[0].landmark]).flatten()
    else:
        lm = np.zeros(478*3)
    landmarks = torch.tensor(lm, dtype=torch.float32).unsqueeze(0)
    
    with torch.no_grad():
        logits = model(pixel_values, landmarks)
        pred = logits.argmax(1).item()
    return pred