wetdog/MOSA-Net_plus

1import os2import torch3import argparse4import numpy as np 5from transformers import AutoFeatureExtractor, WhisperModel6 7import torchaudio8import torch.nn as nn9import torch.nn.functional as F10 11import speechbrain12import librosa13 14from subprocess import CalledProcessError, run15 16#openai whispers load audio17SAMPLE_RATE=1600018def denorm(input_x):19    input_x = input_x*(5-0) + 020    return input_x21    22def load_audio(file: str, sr: int = SAMPLE_RATE):23    """24    Open an audio file and read as mono waveform, resampling as necessary25 26    Parameters27    ----------28    file: str29        The audio file to open30 31    sr: int32        The sample rate to resample the audio if necessary33 34    Returns35    -------36    A NumPy array containing the audio waveform, in float32 dtype.37    """38 39    # This launches a subprocess to decode audio while down-mixing40    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.41    # fmt: off42    cmd = [43        "ffmpeg",44        "-nostdin",45        "-threads", "0",46        "-i", file,47        "-f", "s16le",48        "-ac", "1",49        "-acodec", "pcm_s16le",50        "-ar", str(sr),51        "-"52    ]53    # fmt: on54    try:55        out = run(cmd, capture_output=True, check=True).stdout56    except CalledProcessError as e:57        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e58 59    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.060 61class MosPredictor(nn.Module):62    63    def __init__(self):64        super().__init__()65        66        self.mean_net_conv = nn.Sequential(67            nn.Conv2d(in_channels = 1, out_channels = 16, kernel_size = (3,3), padding = (1,1)),68            nn.Conv2d(in_channels = 16, out_channels = 16, kernel_size = (3,3), padding = (1,1)),69            nn.Conv2d(in_channels = 16, out_channels = 16, kernel_size = (3,3), padding = (1,1), stride=(1,3)),70            nn.Dropout(0.3),71            nn.BatchNorm2d(16),72            nn.ReLU(),73            nn.Conv2d(in_channels = 16, out_channels = 32, kernel_size = (3,3), padding = (1,1)),74            nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = (3,3), padding = (1,1)),75            nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = (3,3), padding = (1,1), stride=(1,3)),76            nn.Dropout(0.3),77            nn.BatchNorm2d(32),78            nn.ReLU(),79            nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = (3,3), padding = (1,1)),80            nn.Conv2d(in_channels = 64, out_channels = 64, kernel_size = (3,3), padding = (1,1)),81            nn.Conv2d(in_channels = 64, out_channels = 64, kernel_size = (3,3), padding = (1,1), stride=(1,3)),82            nn.Dropout(0.3),83            nn.BatchNorm2d(64),84            nn.ReLU(),85            nn.Conv2d(in_channels = 64, out_channels = 128, kernel_size = (3,3), padding = (1,1)),86            nn.Conv2d(in_channels = 128, out_channels = 128, kernel_size = (3,3), padding = (1,1)),87            nn.Conv2d(in_channels = 128, out_channels = 128, kernel_size = (3,3), padding = (1,1), stride=(1,3)),88            nn.Dropout(0.3),89            nn.BatchNorm2d(128),90            nn.ReLU())91        92        self.relu_ = nn.ReLU()93        self.sigmoid_ = nn.Sigmoid()94        95        self.ssl_features = 128096        self.dim_layer = nn.Linear(self.ssl_features, 512)97 98        self.mean_net_rnn = nn.LSTM(input_size = 512, hidden_size = 128, num_layers = 1, batch_first = True, bidirectional = True)99        self.mean_net_dnn = nn.Sequential(100            nn.Linear(256, 128),101            nn.ReLU(),102            nn.Dropout(0.3),103        )        104 105        self.sinc = speechbrain.nnet.CNN.SincConv(in_channels=1, out_channels=257, kernel_size=251, stride=256, sample_rate=16000)106        self.att_output_layer_quality = nn.MultiheadAttention(128, num_heads=8)                107        self.output_layer_quality = nn.Linear(128, 1)108        self.qualaverage_score = nn.AdaptiveAvgPool1d(1)  109     110        self.att_output_layer_intell = nn.MultiheadAttention(128, num_heads=8)           111        self.output_layer_intell = nn.Linear(128, 1)112        self.intellaverage_score = nn.AdaptiveAvgPool1d(1)  113                       114        self.att_output_layer_stoi= nn.MultiheadAttention(128, num_heads=8)          115        self.output_layer_stoi = nn.Linear(128, 1)        116        self.stoiaverage_score = nn.AdaptiveAvgPool1d(1) 117 118    def new_method(self):119        self.sin_conv 120                121    def forward(self, wav, lps, whisper):122        #SSL Features123        wav_ = wav.squeeze(1)  ## [batches, audio_len]124        ssl_feat_red = self.dim_layer(whisper.squeeze(1))125        ssl_feat_red = self.relu_(ssl_feat_red)126 127        #PS Features128        sinc_feat=self.sinc(wav.squeeze(1))129        unsq_sinc =  torch.unsqueeze(sinc_feat, axis=1)130        concat_lps_sinc = torch.cat((lps,unsq_sinc), axis=2)131        cnn_out = self.mean_net_conv(concat_lps_sinc)132        batch = concat_lps_sinc.shape[0]133        time = concat_lps_sinc.shape[2]        134        re_cnn = cnn_out.view((batch, time, 512))135        136        concat_feat = torch.cat((re_cnn,ssl_feat_red), axis=1)137        out_lstm, (h, c) = self.mean_net_rnn(concat_feat)138        out_dense = self.mean_net_dnn(out_lstm) # (batch, seq, 1)       139        140        quality_att, _ = self.att_output_layer_quality (out_dense, out_dense, out_dense) 141        frame_quality = self.output_layer_quality(quality_att)142        frame_quality = self.sigmoid_(frame_quality)   143        quality_utt = self.qualaverage_score(frame_quality.permute(0,2,1))144 145        int_att, _ = self.att_output_layer_intell (out_dense, out_dense, out_dense) 146        frame_int = self.output_layer_intell(int_att)147        frame_int = self.sigmoid_(frame_int)   148        int_utt = self.intellaverage_score(frame_int.permute(0,2,1))149 150                151        return quality_utt.squeeze(1), int_utt.squeeze(1), frame_quality.squeeze(2), frame_int.squeeze(2)152 153