import torch import torch.nn as nn from transformers import Wav2Vec2Processor from transformers.models.wav2vec2.modeling_wav2vec2 import ( Wav2Vec2Model, Wav2Vec2PreTrainedModel, ) import os import librosa import numpy as np 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 # load model from hub device = 'cuda' if torch.cuda.is_available() else "cpu" model_name = 'audeering/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, embeddings: bool = False, ) -> np.ndarray: r"""Predict emotions or extract embeddings from raw audio signal.""" # run through processor to normalize signal # always returns a batch, so we just get the first entry # then we put it on the device y = processor(x, sampling_rate=sampling_rate) y = y['input_values'][0] y = torch.from_numpy(y).to(device) y = y.unsqueeze(0) # run through model with torch.no_grad(): y = model(y)[0 if embeddings else 1] # convert to numpy y = y.detach().cpu().numpy() return y # # # def disp(rootpath, wavname): # wav, sr = librosa.load(f"{rootpath}/{wavname}", 16000) # display(ipd.Audio(wav, rate=sr)) rootpath = "dataset" embs = [] wavnames = [] def extract_dir(path): rootpath = path for idx, wavname in enumerate(os.listdir(rootpath)): wav, sr =librosa.load(f"{rootpath}/{wavname}", sr=16000) emb = process_func(np.expand_dims(wav, 0), sr, embeddings=True) embs.append(emb) wavnames.append(wavname) np.save(f"{rootpath}/{wavname}.emo.npy", emb.squeeze(0)) print(idx, wavname) def extract_wav(path): wav, sr = librosa.load(path, sr=16000) emb = process_func(np.expand_dims(wav, 0), sr, embeddings=True) return emb if __name__ == '__main__': # for spk in ["serena", "koni", "nyaru","shanoa", "mana"]: for spk in ["dubbingx"]: extract_dir(f"dataset/{spk}")