# Model Card: NextViT-Small for Facial Expression Recognition (FER)

This model card provides comprehensive details, training configurations, and performance benchmarks for the fine-tuned **NextViT-Small** model adapted for Facial Expression Recognition (FER).

## Model Details

* **Model Name:** NextViT-Small FER
* **Architecture:** NextViT-Small (`nextvit_small`), a hybrid CNN-Transformer vision model designed for high-performance and hardware-efficient inference on edge devices.
* **Developer:** Fine-tuned locally.
* **Task:** Multi-class image classification for Facial Expression Recognition (7 classes).
* **Number of Classes:** 7 (`angry`, `disgust`, `fear`, `happy`, `neutral`, `sad`, `surprise`).
* **Parameters:** ~28 million parameters (~112 MB for float32 weights-only; checkpoint size is ~381.75 MB including optimizer state).
* **Base Model Weight Source:** Pre-trained on ImageNet (`ckpt5.pth`).

## Training Configuration & Hyperparameters

The model was fine-tuned for 50 epochs on a custom facial expression dataset using the following training parameters:

| Hyperparameter | Value | Description |
| :--- | :--- | :--- |
| **Optimizer** | AdamW | Weight decay set to `0.01` |
| **Learning Rate (Base)** | `2e-4` | Cosine learning rate scheduler with warmup |
| **Epochs** | `50` | Total training epochs |
| **Batch Size** | `128` | Per-GPU batch size (effective batch size of 256 on 2 GPUs) |
| **Input Image Size** | `224 x 224` | Bilinear resize and normalized to ImageNet statistics |
| **Data Augmentations** | RandAugment (`rand-m5-mstd0.5-inc2`), Random Erasing (`reprob 0.1`) | Mixup and Cutmix were disabled (`mixup 0`, `cutmix 0`) |
| **Distributed Training** | PyTorch Distributed Data Parallel (DDP) | 2 GPUs (Distributed evaluation enabled, NCCL backend) |
| **Total Training Time** | `0:41:41` (41m 41s) | Training completed on a dual-GPU node |

---

## Evaluation Metrics & Performance

### 1. General Test Set Performance (12,957 Images)
At the end of training (Epoch 50), the model achieved the following performance on the full validation/test dataset:
* **Top-1 Accuracy:** **82.20%**
* **Top-5 Accuracy:** **99.04%**
* **Final Test Loss:** **0.739**

### 2. Detailed Test Split Performance (2,463 Images)
When evaluated on a dedicated test split of 2,463 samples, NextViT-Small achieved:
* **Overall Accuracy:** **83.88%**
* **Precision (Macro):** **75.92%**
* **Precision (Weighted):** **83.82%**
* **Recall (Macro):** **73.93%**
* **Recall (Weighted):** **83.88%**
* **F1-Score (Macro):** **74.61%**
* **F1-Score (Weighted):** **83.78%**

#### Per-Class Performance Breakdown:
| Class | Class ID | Precision | Recall | F1-Score | Support |
| :--- | :---: | :---: | :---: | :---: | :---: |
| **happy** | 3 | 95.51% | 93.37% | 94.43% | 935 |
| **surprise** | 6 | 83.33% | 84.34% | 83.83% | 249 |
| **angry** | 0 | 79.38% | 83.01% | 81.15% | 153 |
| **neutral** | 4 | 77.84% | 81.89% | 79.81% | 519 |
| **sad** | 5 | 78.67% | 80.39% | 79.52% | 413 |
| **disgust** | 1 | 58.78% | 56.62% | 57.68% | 136 |
| **fear** | 2 | 57.89% | 37.93% | 45.83% | 58 |

---

### 3. Comparison with ViT-Small Baseline
NextViT-Small shows significant improvements over a standard vision transformer (ViT-Small) baseline evaluated on the same 2,463 image test split:

* **ViT-Small Accuracy:** **80.31%** (Macro F1-Score: **71.39%**)
* **NextViT-Small Accuracy:** **83.88%** (Macro F1-Score: **74.61%**)
* **Performance Gain:** **+3.57% Accuracy** improvement with NextViT-Small.

NextViT-Small demonstrates much stronger capability on classes with lower support (e.g. `angry`, `disgust`) due to its hybrid convolutional-attention structure, which excels in learning fine-grained local facial details.

---

## How to Use

To run inference using the trained NextViT-Small model, follow this example snippet:

```python
import os
import sys
import torch
from PIL import Image
from torchvision import transforms

# 1. Path setup: Add Next-ViT classification folder to path so timm can register nextvit_small
NEXTVIT_PATH = "/path/to/Next-ViT/classification" # Update as appropriate
sys.path.append(NEXTVIT_PATH)
import nextvit # registers 'nextvit_small' with timm
import timm

# 2. Load the model architecture and pretrained weights
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = timm.create_model("nextvit_small", num_classes=7)

checkpoint_path = "checkpoint_best.pth" # Update path to best checkpoint
checkpoint = torch.load(checkpoint_path, map_location=device)

if "model" in checkpoint:
    model.load_state_dict(checkpoint["model"])
else:
    model.load_state_dict(checkpoint)

model.to(device)
model.eval()

# 3. Define the image preprocessing transforms
image_processor = 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)),
])

# 4. Run inference
class_names = ["angry", "disgust", "fear", "happy", "neutral", "sad", "surprise"]

def predict_expression(image_path):
    img = Image.open(image_path).convert("RGB")
    tensor = image_processor(img).unsqueeze(0).to(device)
    
    with torch.no_grad():
        outputs = model(tensor)
        probs = torch.nn.functional.softmax(outputs, dim=-1)
        pred_idx = torch.argmax(probs, dim=-1).item()
        
    return class_names[pred_idx], probs[0][pred_idx].item()

# Example Prediction
# expression, confidence = predict_expression("test_face.jpg")
# print(f"Predicted Expression: {expression} ({confidence:.2%})")
```

---

## Intended Use & Limitations

* **Intended Use:** Facial Expression Recognition in user-facing applications (e.g. human-computer interaction, affective computing, user experience testing).
* **Limitations:** Performance on expressions with very low support (such as `fear` or `disgust`) remains relatively lower compared to well-represented expressions (like `happy` and `neutral`). Users should be cautious when deploying in scenarios sensitive to false negatives on these classes.