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