app.py

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('''<table style="width:500px"><tr><th colspan=5 >CCC MSP Podcast v1.7</th></tr>
  <tr> <td> </td><td>Arousal</td> <td>Dominance</td> <td>Valence</td> <td> Associated Paper </td> </tr>
  <tr> <td> <a href="https://huggingface.co/audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim">Wav2Vec2</a></td><td>0.744</td><td>0.655</td><td> 0.638 </td><td> <a href="https://arxiv.org/abs/2203.07378">arXiv</a> </td> </tr>
  <tr> <td> <a href="https://huggingface.co/dkounadis/wav2small">Wav2Small Teacher</a></td><td> 0.762 </td> <td> 0.684 </td><td> 0.676 </td><td> <a href="https://arxiv.org/abs/2408.13920">arXiv</a> </td> </tr>
</table>
''')

# Launch the Gradio application
if __name__ == "__main__":
    demo.launch(share=False)