---
library_name: timm
tags:
- vision
- image-classification
- facial-expression-recognition
- vit
- vision-transformer
- fer
- ferplus
datasets:
- ferplus
metrics:
- accuracy
- f1
- precision
- recall
pipeline_tag: image-classification
---

# 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:

| Hyperparameter | Value |
| :--- | :--- |
| **Epochs** | 50 |
| **Batch Size** | 128 |
| **Base Learning Rate** | 5e-4 |
| **Weight Decay** | 0.01 |
| **Optimizer** | AdamW |
| **Scheduler** | Cosine Annealing (with 5-epoch warmup starting at 1e-6, min LR 1e-5) |
| **Image Resolution** | 224 x 224 (Bicubic Interpolation) |
| **Stochastic Depth (Drop Path)** | 0.1 |
| **Precision** | Mixed Precision (AMP) |
| **Data Augmentations** | RandAugment, Random Erasing (p=0.1) |
| **Loss Function** | Weighted Cross-Entropy Loss |
| **Class Weights** | `[1.560, 3.737, 2.242, 0.541, 0.527, 0.970, 1.149]` (mapping to classes 0 to 6) |

---

## 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:
| Class ID | Class Name | Precision | Recall | F1-Score | Support |
| :---: | :--- | :---: | :---: | :---: | :---: |
| 3 | Happy | 0.9247 | 0.9345 | 0.9295 | 473 |
| 6 | Surprise | 0.8636 | 0.7661 | 0.8120 | 124 |
| 4 | Neutral | 0.7943 | 0.8145 | 0.8043 | 275 |
| 0 | Angry | 0.7412 | 0.8400 | 0.7875 | 75 |
| 5 | Sad | 0.7758 | 0.7711 | 0.7734 | 166 |
| 2 | Fear | 0.8000 | 0.7273 | 0.7619 | 33 |
| 1 | Disgust | 0.6250 | 0.5833 | 0.6034 | 60 |

---

### 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:
| Class ID | Class Name | Precision | Recall | F1-Score | Support |
| :---: | :--- | :---: | :---: | :---: | :---: |
| 3 | Happy | 0.9187 | 0.9378 | 0.9281 | 482 |
| 6 | Surprise | 0.8793 | 0.7969 | 0.8361 | 128 |
| 0 | Angry | 0.8158 | 0.8493 | 0.8322 | 73 |
| 5 | Sad | 0.8688 | 0.7277 | 0.7920 | 191 |
| 4 | Neutral | 0.7240 | 0.8707 | 0.7906 | 232 |
| 2 | Fear | 0.6471 | 0.5000 | 0.5641 | 22 |
| 1 | Disgust | 0.6066 | 0.5068 | 0.5522 | 73 |

---

## How to Use

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

```python
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%})")
```