ViTFERPP_vit_small_patch16_224.augreg_in1k

ViT-Small for Facial Expression Recognition (FER++)

This repository contains a Vision Transformer (ViT-Small) model fine-tuned for Facial Expression Recognition (FER).

The model is initialized from the pretrained checkpoint vit_small_patch16_224.augreg_in1k from the timm library and fine-tuned on a 7-class emotion classification dataset (aligned with FER++ taxonomy).

Model Details

• Model Architecture: Vision Transformer (ViT-Small)
• Pretrained Base: vit_small_patch16_224.augreg_in1k
• Number of Parameters: ~22M
• Task: 7-class Facial Expression Classification
• Classes:
  • 0: Angry
  • 1: Disgust
  • 2: Fear
  • 3: Happy
  • 4: Neutral
  • 5: Sad
  • 6: Surprise

Training Parameters & Hyperparameters

The model was fine-tuned using PyTorch and the timm library with the following setup:

---

Evaluation Results

Validation Set Performance

• Accuracy: 83.83%
• Precision (Macro): 0.7892
• Recall (Macro): 0.7767
• F1-Score (Macro): 0.7817
• F1-Score (Weighted): 0.8378

Per-Class Validation Metrics:

---

Test Set Performance

• Accuracy: 83.68%
• Precision (Macro): 0.7800
• Recall (Macro): 0.7413
• F1-Score (Macro): 0.7565
• F1-Score (Weighted): 0.8348

Per-Class Test Metrics:

---

How to Use

Here is how you can load the model and run inference on an image using PyTorch and timm:

import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from timm import create_model

# Define class mapping
class_mapping = {
    0: "Angry",
    1: "Disgust",
    2: "Fear",
    3: "Happy",
    4: "Neutral",
    5: "Sad",
    6: "Surprise"
}

# 1. Define image preprocessing matching validation/test set config
transform = transforms.Compose([
    transforms.Resize((224, 224), interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

# 2. Instantiate and build the model structure
model = create_model('vit_small_patch16_224.augreg_in1k', pretrained=False, num_classes=7)

# 3. Load weight state dictionary
state_dict = torch.load('vitsmall2.pth', map_location='cpu')

# Clean state_dict if it was saved using DataParallel (remove 'module.' prefix)
if list(state_dict.keys())[0].startswith('module.'):
    state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}

model.load_state_dict(state_dict)
model.eval()

# 4. Perform Inference
image_path = "path_to_facial_image.jpg"
image = Image.open(image_path).convert('RGB')
input_tensor = transform(image).unsqueeze(0)  # Add batch dimension

with torch.no_grad():
    outputs = model(input_tensor)
    probabilities = torch.softmax(outputs, dim=1)[0]
    predicted_class_id = torch.argmax(probabilities).item()

print(f"Predicted emotion: {class_mapping[predicted_class_id]} ({probabilities[predicted_class_id].item():.2%})")