import gradio as gr import torch.nn as nn import audresample import matplotlib.pyplot as plt from matplotlib import colors as mcolors import torch import librosa import numpy as np import types from transformers import AutoModelForAudioClassification from transformers.models.wav2vec2.modeling_wav2vec2 import (Wav2Vec2Model, Wav2Vec2PreTrainedModel) plt.style.use('seaborn-v0_8-whitegrid') class ADV(nn.Module): def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.out_proj = nn.Linear(config.hidden_size, config.num_labels) def forward(self, x): x = self.dense(x) x = torch.tanh(x) return self.out_proj(x) class Dawn(Wav2Vec2PreTrainedModel): r"""https://arxiv.org/abs/2203.07378""" def __init__(self, config): super().__init__(config) self.wav2vec2 = Wav2Vec2Model(config) self.classifier = ADV(config) def forward(self, x): x -= x.mean(1, keepdim=True) variance = (x * x).mean(1, keepdim=True) + 1e-7 x = self.wav2vec2(x / variance.sqrt()) return self.classifier(x.last_hidden_state.mean(1)) def _forward(self, x): '''x: (batch, audio-samples-16KHz)''' x = (x + self.config.mean) / self.config.std # sgn x = self.ssl_model(x, attention_mask=None).last_hidden_state # pool h = self.pool_model.sap_linear(x).tanh() w = torch.matmul(h, self.pool_model.attention).softmax(1) mu = (x * w).sum(1) x = torch.cat( [ mu, ((x * x * w).sum(1) - mu * mu).clamp(min=1e-7).sqrt() ], 1) return self.ser_model(x) # WavLM device = 'cpu' base = AutoModelForAudioClassification.from_pretrained( '3loi/SER-Odyssey-Baseline-WavLM-Multi-Attributes', trust_remote_code=True).to(device).eval() base.forward = types.MethodType(_forward, base) # Wav2Vec2 dawn = Dawn.from_pretrained( 'audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim' ).to(device).eval() # Wav2Small import torch import numpy as np import torch.nn.functional as F import librosa from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2PreTrainedModel, Wav2Vec2Model from torch import nn from transformers import PretrainedConfig def _prenorm(x, attention_mask=None): '''mean/var''' if attention_mask is not None: N = attention_mask.sum(1, keepdim=True) # here attn msk is unprocessed just the original input x -= x.sum(1, keepdim=True) / N var = (x * x).sum(1, keepdim=True) / N else: x -= x.mean(1, keepdim=True) # mean is an onnx operator reducemean saves some ops compared to casting integer N to float and the div var = (x * x).mean(1, keepdim=True) return x / torch.sqrt(var + 1e-7) class Spectrogram(nn.Module): def __init__(self, n_fft=64, # num cols of DFT n_time=64, # num rows of DFT matrix hop_length=32, freeze_parameters=True): super().__init__() fft_window = librosa.filters.get_window('hann', n_time, fftbins=True) fft_window = librosa.util.pad_center(fft_window, size=n_time) out_channels = n_fft // 2 + 1 (x, y) = np.meshgrid(np.arange(n_time), np.arange(n_fft)) omega = np.exp(-2 * np.pi * 1j / n_time) dft_matrix = np.power(omega, x * y) # (n_fft, n_time) dft_matrix = dft_matrix * fft_window[None, :] dft_matrix = dft_matrix[0 : out_channels, :] dft_matrix = dft_matrix[:, None, :] # ---- Assymetric DFT Non Square self.conv_real = nn.Conv1d(1, out_channels, n_fft, stride=hop_length, padding=0, bias=False) self.conv_imag = nn.Conv1d(1, out_channels, n_fft, stride=hop_length, padding=0, bias=False) self.conv_real.weight.data = torch.tensor(np.real(dft_matrix), dtype=self.conv_real.weight.dtype).to(self.conv_real.weight.device) self.conv_imag.weight.data = torch.tensor(np.imag(dft_matrix), dtype=self.conv_imag.weight.dtype).to(self.conv_imag.weight.device) if freeze_parameters: for param in self.parameters(): param.requires_grad = False def forward(self, input): x = input[:, None, :] real = self.conv_real(x) imag = self.conv_imag(x) return real ** 2 + imag ** 2 # bs, mel, time-frames class LogmelFilterBank(nn.Module): def __init__(self, sr=16000, n_fft=64, n_mels=26, # maxpool fmin=0.0, freeze_parameters=True): super().__init__() fmax = sr//2 W2 = librosa.filters.mel(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax).T self.register_buffer('melW', torch.Tensor(W2)) self.register_buffer('amin', torch.Tensor([1e-10])) def forward(self, x): x = torch.matmul(x[:, None, :, :].transpose(2, 3), self.melW) # changes melf not num frames x = torch.where(x > self.amin, x, self.amin) # not in place x = 10 * torch.log10(x) return x def length_after_conv_layer(_length, k=None, pad=None, stride=None): return torch.floor( (_length + 2*pad - k) / stride + 1 ) class Conv(nn.Module): def __init__(self, c_in, c_out, k=3, stride=1, padding=1): super().__init__() self.conv = nn.Conv2d(c_in, c_out, k, stride=stride, padding=padding, bias=False) self.norm = nn.BatchNorm2d(c_out) def forward(self, x): x = self.conv(x) x = self.norm(x) return torch.relu_(x) class Vgg7(nn.Module): def __init__(self): super().__init__() self.l1 = Conv( 1, 13) self.l2 = Conv(13, 13) self.l3 = Conv(13, 13) self.maxpool_A = nn.MaxPool2d(3, stride=2, padding=1) self.l4 = Conv(13, 13) self.l5 = Conv(13, 13) self.l6 = Conv(13, 13) self.l7 = Conv(13, 13) self.lin = nn.Conv2d(13, 13, 1, padding=0, stride=1) self.sof = nn.Conv2d(13, 13, 1, padding=0, stride=1) # pool time - reshape mel into channels after pooling self.spectrogram_extractor = Spectrogram() self.logmel_extractor = LogmelFilterBank() def final_length(self, L): conv_kernel = [64, 3] # [nfft, maxpool] conv_stride = [32, 2] # [hop_len, maxpool_stride] # consider only layers of stride > 1 conv_pad = [0, 1] # [pad_stft, pad_maxpool] for k, stride, pad in zip(conv_kernel, conv_stride, conv_pad): L = length_after_conv_layer(L, k=k, stride=stride, pad=pad) return L def final_attention_mask(self, feature_vector_length, attention_mask=None): non_padded_lengths = attention_mask.sum(1) out_lengths = self.final_length(non_padded_lengths) # how can non_padded_lengths get exact 0 here DOES IT MEAN ATTNMASK WAS NOT FILLED? out_lengths = out_lengths.to(torch.long) bs, _ = attention_mask.shape attention_mask = torch.ones((bs, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device) for b, _len in enumerate(out_lengths): attention_mask[b, _len:] = 0 return attention_mask def forward(self, x, attention_mask=None): x = _prenorm(x, attention_mask=attention_mask) x = self.spectrogram_extractor(x) x = self.logmel_extractor(x) x = self.l1(x) x = self.l2(x) x = self.l3(x) x = self.maxpool_A(x) # reshape here? so these conv will have large kernel x = self.l4(x) x = self.l5(x) x = self.l6(x) x = self.l7(x) if attention_mask is not None: bs, _, t, _ = x.shape a = self.final_attention_mask(feature_vector_length=t, attention_mask=attention_mask)[:, None, :, None] #print(a.shape, x.shape, '\n\n\n\n') x = torch.masked_fill(x, a < 1, 0) # mask also affects lin !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! x = self.lin(x) * ( self.sof(x) -10000. * torch.logical_not(a) ).softmax(2) else: x = self.lin(x) * self.sof(x).softmax(2) x = x.sum(2) # bs, ch, time-frames, HALF_MEL -> bs, ch, HALF_MEL # -- xT = x.transpose(1,2) x = torch.cat([x, torch.bmm(x, xT), # corr (chxmel) x (melxCH) # torch.bmm(x, x), # corr ch * ch # torch.bmm(xT, xT) # corr mel * mel ], 2) # -- return x.reshape(-1, 338) class Wav2SmallConfig(PretrainedConfig): model_type = "wav2vec2" def __init__(self, **kwargs): super().__init__(**kwargs) self.half_mel = 13 self.n_fft = 64 self.n_time = 64 self.hidden = 2 * self.half_mel * self.half_mel self.hop = self.n_time // 2 class Wav2Small(Wav2Vec2PreTrainedModel): def __init__(self, config): super().__init__(config) self.vgg7 = Vgg7() self.adv = nn.Linear(config.hidden, 3) # 0=arousal, 1=dominance, 2=valence def forward(self, x, attention_mask=None): x = self.vgg7(x, attention_mask=attention_mask) return self.adv(x) def _ccc(x, y): '''if len(x) = len(y) = 1 we have 0/0 as a&b can both be negative we should add 1e-7 to denominator protecting sign of denominator to find sign of denominator and add 1e-7 if sgn>=0 or -1e-7 if sgn<0''' mean_y = y.mean() mean_x = x.mean() a = x - mean_x b = y - mean_y L = (mean_x - mean_y).abs() * .1 * x.shape[0] #print(L / ((mean_x - mean_y) **2 * x.shape[0])) numerator = torch.dot(a, b) # L term if both a,b scalars dissallows 0 numerator [OFFICIAL CCC HAS L ONLY IN D] denominator = torch.dot(a, a) + torch.dot(b, b) + L # if both a,b are equalscalars then the dots are all zero and ccc=1 denominator = torch.where(denominator.sign() < 0, denominator - 1e-7, denominator + 1e-7) ccc = numerator / denominator return -ccc #+ F.l1_loss(a, b) wav2small = Wav2Small.from_pretrained('audeering/wav2small').to(device).eval() # Error figure for the first plot fig_error, ax = plt.subplots(figsize=(8, 6)) error_message = "Error: No .wav or Mic. audio provided." ax.text(0.5, 0.5, error_message, ha='center', va='center', fontsize=24, color='gray', fontweight='bold', transform=ax.transAxes) ax.set_xticks([]) ax.set_yticks([]) ax.set_xticklabels([]) ax.set_yticklabels([]) ax.set_frame_on(True) ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) ax.spines['bottom'].set_visible(False) ax.spines['left'].set_visible(False) def process_audio(audio_filepath): if audio_filepath is None: return fig_error, fig_error waveform, sample_rate = librosa.load(audio_filepath, sr=None) # Resample audio to 16kHz if the sample rate is different if sample_rate != 16000: resampled_waveform_np = audresample.resample(waveform, sample_rate, 16000) else: resampled_waveform_np = waveform[None, :] x = torch.from_numpy(resampled_waveform_np[:, :64000]).to(torch.float) # only 4s for speed with torch.no_grad(): logits_dawn = dawn(x).cpu().numpy()[0, :] logits_wavlm = base(x).cpu().numpy()[0, :] # 17K params logits_wav2small = wav2small(x).cpu().numpy()[0, :] # --- Plot 1: Wav2Vec2 vs Wav2Small Teacher Outputs --- fig, ax = plt.subplots(figsize=(10, 6)) left_bars_data = logits_dawn.clip(0, 1) right_bars_data = logits_wav2small.clip(0, 1) bar_labels = ['\nArousal', '\nDominance', '\nValence'] y_pos = np.arange(len(bar_labels)) # Define colormaps for each category to ensure distinct colors category_colormaps = [plt.cm.Blues, plt.cm.Greys, plt.cm.Oranges] left_filled_colors = [] right_filled_colors = [] background_colors = [] # Assign specific shades for filled bars and background bars for i, cmap in enumerate(category_colormaps): left_filled_colors.append(cmap(0.74)) right_filled_colors.append(cmap(0.64)) background_colors.append(cmap(0.1)) # Plot transparent background bars for i in range(len(bar_labels)): ax.barh(y_pos[i], -1, color=background_colors[i], alpha=0.3, height=0.6) ax.barh(y_pos[i], 1, color=background_colors[i], alpha=0.3, height=0.6) # Plot the filled bars for actual data for i in range(len(bar_labels)): ax.barh(y_pos[i], -left_bars_data[i], color=left_filled_colors[i], alpha=1, height=0.6) ax.barh(y_pos[i], right_bars_data[i], color=right_filled_colors[i], alpha=1, height=0.6) # Add a central vertical axis divider ax.axvline(0, color='black', linewidth=0.8, linestyle='--') # Set x-axis limits and y-axis ticks/labels ax.set_xlim(-1, 1) ax.set_yticks(y_pos) ax.set_yticklabels(bar_labels, fontsize=12) # Custom formatter for x-axis to show absolute percentage values def abs_tick_formatter(x, pos): return f'{int(abs(x) * 100)}%' ax.xaxis.set_major_formatter(plt.FuncFormatter(abs_tick_formatter)) # Set plot title and x-axis label ax.set_title('', fontsize=16, pad=20) ax.set_xlabel('Wav2Vev2 (Dawn) Wav2Small (17K param.)', fontsize=12) # Remove top, right, and left spines for a cleaner look ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) ax.spines['left'].set_visible(False) # Add annotations (percentage values) to the filled bars for i in range(len(bar_labels)): ax.text(-left_bars_data[i] - 0.05, y_pos[i], f'{int(left_bars_data[i] * 100)}%', va='center', ha='right', color=left_filled_colors[i], fontweight='bold') ax.text(right_bars_data[i] + 0.05, y_pos[i], f'{int(right_bars_data[i] * 100)}%', va='center', ha='left', color=right_filled_colors[i], fontweight='bold') # -- PLOT 2 : WavLM / Wav2Small Teacher fig_2, ax_2 = plt.subplots(figsize=(10, 6)) left_bars_data = logits_wavlm.clip(0, 1) right_bars_data = (.5 * logits_dawn + .5 * logits_wavlm).clip(0, 1) bar_labels = ['\nArousal', '\nDominance', '\nValence'] y_pos = np.arange(len(bar_labels)) # Define colormaps for each category to ensure distinct colors category_colormaps = [plt.cm.Blues, plt.cm.Greys, plt.cm.Oranges] left_filled_colors = [] right_filled_colors = [] background_colors = [] # Assign specific shades for filled bars and background bars for i, cmap in enumerate(category_colormaps): left_filled_colors.append(cmap(0.74)) right_filled_colors.append(cmap(0.64)) background_colors.append(cmap(0.1)) # Plot transparent background bars for i in range(len(bar_labels)): ax_2.barh(y_pos[i], -1, color=background_colors[i], alpha=0.3, height=0.6) ax_2.barh(y_pos[i], 1, color=background_colors[i], alpha=0.3, height=0.6) # Plot the filled bars for actual data for i in range(len(bar_labels)): ax_2.barh(y_pos[i], -left_bars_data[i], color=left_filled_colors[i], alpha=1, height=0.6) ax_2.barh(y_pos[i], right_bars_data[i], color=right_filled_colors[i], alpha=1, height=0.6) # Add a central vertical axis divider ax_2.axvline(0, color='black', linewidth=0.8, linestyle='--') # Set x-axis limits and y-axis ticks/labels ax_2.set_xlim(-1, 1) ax_2.set_yticks(y_pos) ax_2.set_yticklabels(bar_labels, fontsize=12) # Custom formatter for x-axis to show absolute percentage values def abs_tick_formatter(x, pos): return f'{int(abs(x) * 100)}%' ax_2.xaxis.set_major_formatter(plt.FuncFormatter(abs_tick_formatter)) ax_2.set_title('', fontsize=16, pad=20) ax_2.set_xlabel('WavLM (Baseline) Wav2Small Teacher (0.4B param.)', fontsize=12) ax_2.spines['top'].set_visible(False) ax_2.spines['right'].set_visible(False) ax_2.spines['left'].set_visible(False) # Add annotations (percentage values) to the filled bars for i in range(len(bar_labels)): ax_2.text(-left_bars_data[i] - 0.05, y_pos[i], f'{int(left_bars_data[i] * 100)}%', va='center', ha='right', color=left_filled_colors[i], fontweight='bold') ax_2.text(right_bars_data[i] + 0.05, y_pos[i], f'{int(right_bars_data[i] * 100)}%', va='center', ha='left', color=right_filled_colors[i], fontweight='bold') return fig, fig_2 iface = gr.Interface( fn=process_audio, inputs=gr.Audio( sources=["microphone", "upload"], type="filepath", # Input type is file path label='' ), outputs=[ gr.Plot(label="Wav2Vec2 vs Wav2Small (17K params) Plot"), # First plot output gr.Plot(label="WavLM vs Wav2Small Teacher Plot"), # Second plot output ], title='', description='', flagging_mode="never", # Disables flagging feature examples=[ "female-46-neutral.wav", "female-20-happy.wav", "male-60-angry.wav", "male-27-sad.wav", ], css="footer {visibility: hidden}" # Hides the Gradio footer ) # Gradio Blocks for tabbed interface with gr.Blocks() as demo: # First tab for the existing Arousal/Dominance/Valence plots with gr.Tab(label="Arousal / Dominance / Valence"): iface.render() # Second tab for CCC (Concordance Correlation Coefficient) information with gr.Tab(label="CCC"): gr.Markdown('''
CCC MSP Podcast v1.7
Arousal Dominance Valence Associated Paper
Wav2Vec20.7440.655 0.638 arXiv
Wav2Small Teacher 0.762 0.684 0.676 arXiv
''') # Launch the Gradio application if __name__ == "__main__": demo.launch(share=False)