Roi Feng
new 2.2
2e9bf0c
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)
# 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 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)