File size: 4,744 Bytes
75c1b11 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | ---
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%})")
```
|