| import librosa |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from torch.utils.data import Dataset |
| from torch.utils.data import Dataset |
| from transformers import Wav2Vec2Processor |
| from transformers.models.wav2vec2.modeling_wav2vec2 import ( |
| Wav2Vec2Model, |
| Wav2Vec2PreTrainedModel, |
| ) |
|
|
| from config import config |
|
|
|
|
| class RegressionHead(nn.Module): |
| r"""Classification head.""" |
|
|
| def __init__(self, config): |
| super().__init__() |
|
|
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.dropout = nn.Dropout(config.final_dropout) |
| self.out_proj = nn.Linear(config.hidden_size, config.num_labels) |
|
|
| def forward(self, features, **kwargs): |
| x = features |
| x = self.dropout(x) |
| x = self.dense(x) |
| x = torch.tanh(x) |
| x = self.dropout(x) |
| x = self.out_proj(x) |
|
|
| return x |
|
|
|
|
| class EmotionModel(Wav2Vec2PreTrainedModel): |
| r"""Speech emotion classifier.""" |
|
|
| def __init__(self, config): |
| super().__init__(config) |
|
|
| self.config = config |
| self.wav2vec2 = Wav2Vec2Model(config) |
| self.classifier = RegressionHead(config) |
| self.init_weights() |
|
|
| def forward( |
| self, |
| input_values, |
| ): |
| outputs = self.wav2vec2(input_values) |
| hidden_states = outputs[0] |
| hidden_states = torch.mean(hidden_states, dim=1) |
| logits = self.classifier(hidden_states) |
|
|
| return hidden_states, logits |
|
|
|
|
| class AudioDataset(Dataset): |
| def __init__(self, list_of_wav_files, sr, processor): |
| self.list_of_wav_files = list_of_wav_files |
| self.processor = processor |
| self.sr = sr |
|
|
| def __len__(self): |
| return len(self.list_of_wav_files) |
|
|
| def __getitem__(self, idx): |
| wav_file = self.list_of_wav_files[idx] |
| audio_data, _ = librosa.load(wav_file, sr=self.sr) |
| processed_data = self.processor(audio_data, sampling_rate=self.sr)[ |
| "input_values" |
| ][0] |
| return torch.from_numpy(processed_data) |
|
|
|
|
| device = config.emo_gen_config.device |
| model_name = "./emotional/wav2vec2-large-robust-12-ft-emotion-msp-dim" |
| processor = Wav2Vec2Processor.from_pretrained(model_name) |
| model = EmotionModel.from_pretrained(model_name).to(device) |
|
|
|
|
| def process_func( |
| x: np.ndarray, |
| sampling_rate: int, |
| model: EmotionModel, |
| processor: Wav2Vec2Processor, |
| device: str, |
| embeddings: bool = False, |
| ) -> np.ndarray: |
| r"""Predict emotions or extract embeddings from raw audio signal.""" |
| model = model.to(device) |
| y = processor(x, sampling_rate=sampling_rate) |
| y = y["input_values"][0] |
| y = torch.from_numpy(y).unsqueeze(0).to(device) |
|
|
| |
| with torch.no_grad(): |
| y = model(y)[0 if embeddings else 1] |
|
|
| |
| y = y.detach().cpu().numpy() |
|
|
| return y |
|
|
|
|
| def get_emo(path): |
| wav, sr = librosa.load(path, 16000) |
| return process_func( |
| np.expand_dims(wav, 0).astype(np.float64), |
| sr, |
| model, |
| processor, |
| device, |
| embeddings=True, |
| ).squeeze(0) |
|
|
