Views
No views yet
ViT Base (86M parameters)
↓
CLS Token Output (768-dim)
↓
LayerNorm + Dropout
↓
Linear (768 → 512) + GELU + Dropout
↓
Linear (512 → 128) + GELU + Dropout
↓
Linear (128 → 2) + Tanh
↓
[Valence, Arousal] ∈ [-1, 1]²pip install torch transformers librosa numpy pillow1import torch
2from transformers import ViTModel
3import torch.nn as nn
4
5class ViTForEmotionRegression(nn.Module):
6 def __init__(self, model_name='google/vit-base-patch16-224-in21k', num_emotions=2, dropout=0.1):
7 super().__init__()
8 self.vit = ViTModel.from_pretrained(model_name)
9 hidden_size = self.vit.config.hidden_size
10
11 self.head = nn.Sequential(
12 nn.LayerNorm(hidden_size),
13 nn.Dropout(dropout),
14 nn.Linear(hidden_size, 512),
15 nn.GELU(),
16 nn.Dropout(dropout),
17 nn.Linear(512, 128),
18 nn.GELU(),
19 nn.Dropout(dropout),
20 nn.Linear(128, num_emotions),
21 nn.Tanh()
22 )
23
24 def forward(self, pixel_values):
25 outputs = self.vit(pixel_values)
26 cls_output = outputs.last_hidden_state[:, 0]
27 return self.head(cls_output)
28
29# Load the model
30model = ViTForEmotionRegression()
31model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
32model.eval()1import librosa
2import numpy as np
3from PIL import Image
4import torch
5from torchvision import transforms
6
7def preprocess_audio(audio_path):
8 # Load audio
9 y, sr = librosa.load(audio_path, sr=22050, duration=30)
10
11 # Generate mel spectrogram
12 mel_spec = librosa.feature.melspectrogram(
13 y=y, sr=sr, n_mels=128, hop_length=512, n_fft=2048
14 )
15 mel_db = librosa.power_to_db(mel_spec, ref=np.max)
16
17 # Normalize to 0-255 for RGB conversion
18 mel_normalized = ((mel_db - mel_db.min()) / (mel_db.max() - mel_db.min()) * 255).astype(np.uint8)
19
20 # Convert to RGB image
21 image = Image.fromarray(mel_normalized).convert('RGB')
22 image = image.resize((224, 224))
23
24 # Apply ImageNet normalization
25 transform = transforms.Compose([
26 transforms.ToTensor(),
27 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
28 ])
29
30 return transform(image).unsqueeze(0)
31
32# Process audio
33audio_tensor = preprocess_audio('your_audio.mp3')
34
35# Predict emotions
36with torch.no_grad():
37 predictions = model(audio_tensor)
38 valence, arousal = predictions[0].tolist()
39
40print(f"Valence: {valence:.3f}, Arousal: {arousal:.3f}")1def classify_emotion(valence, arousal):
2 if valence >= 0 and arousal >= 0:
3 return "HAPPY" if valence > arousal else "EXCITED"
4 elif valence >= 0 and arousal < 0:
5 return "CALM" if abs(arousal) > valence else "CONTENT"
6 elif valence < 0 and arousal < 0:
7 return "SAD" if abs(valence) > abs(arousal) else "BORED"
8 else: # valence < 0 and arousal >= 0
9 return "TENSE" if arousal > abs(valence) else "ANGRY" High Arousal
|
Angry Tense Excited
|
Sad -------- + -------- Happy
|
Bored Calm Content
|
Low Arousal1@misc{sentio-vit-emotion,
2 title={Vision Transformer for Audio Emotion Recognition},
3 author={SentioApp Team},
4 year={2025},
5 publisher={HuggingFace}
6}