|
import glob
|
|
import librosa
|
|
import tqdm
|
|
import numpy as np
|
|
import torchaudio
|
|
import torch
|
|
|
|
|
|
import warnings
|
|
|
|
warnings.filterwarnings("ignore")
|
|
|
|
import concurrent.futures
|
|
import glob
|
|
import os
|
|
import librosa
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
import pandas as pd
|
|
from tqdm import tqdm
|
|
|
|
SAMPLING_RATE = 16000
|
|
INPUT_LENGTH = 9.01
|
|
|
|
|
|
class DNSMOSComputer:
|
|
def __init__(
|
|
self, primary_model_path, p808_model_path, device="cuda", device_id=0
|
|
) -> None:
|
|
self.onnx_sess = ort.InferenceSession(
|
|
primary_model_path, providers=["CUDAExecutionProvider"]
|
|
)
|
|
self.p808_onnx_sess = ort.InferenceSession(
|
|
p808_model_path, providers=["CUDAExecutionProvider"]
|
|
)
|
|
self.onnx_sess.set_providers(["CUDAExecutionProvider"], [{"device_id": device_id}])
|
|
self.p808_onnx_sess.set_providers(
|
|
["CUDAExecutionProvider"], [{"device_id": device_id}]
|
|
)
|
|
kwargs = {
|
|
"sample_rate": 16000,
|
|
"hop_length": 160,
|
|
"n_fft": 320 + 1,
|
|
"n_mels": 120,
|
|
"mel_scale": "slaney",
|
|
}
|
|
self.mel_transform = torchaudio.transforms.MelSpectrogram(**kwargs).to(f"cuda:{device_id}")
|
|
|
|
def audio_melspec(
|
|
self, audio, n_mels=120, frame_size=320, hop_length=160, sr=16000, to_db=True
|
|
):
|
|
mel_specgram = self.mel_transform(torch.Tensor(audio).cuda())
|
|
mel_spec = mel_specgram.cpu()
|
|
if to_db:
|
|
mel_spec = (librosa.power_to_db(mel_spec, ref=np.max) + 40) / 40
|
|
return mel_spec.T
|
|
|
|
def get_polyfit_val(self, sig, bak, ovr, is_personalized_MOS):
|
|
if is_personalized_MOS:
|
|
p_ovr = np.poly1d([-0.00533021, 0.005101, 1.18058466, -0.11236046])
|
|
p_sig = np.poly1d([-0.01019296, 0.02751166, 1.19576786, -0.24348726])
|
|
p_bak = np.poly1d([-0.04976499, 0.44276479, -0.1644611, 0.96883132])
|
|
else:
|
|
p_ovr = np.poly1d([-0.06766283, 1.11546468, 0.04602535])
|
|
p_sig = np.poly1d([-0.08397278, 1.22083953, 0.0052439])
|
|
p_bak = np.poly1d([-0.13166888, 1.60915514, -0.39604546])
|
|
sig_poly = p_sig(sig)
|
|
bak_poly = p_bak(bak)
|
|
ovr_poly = p_ovr(ovr)
|
|
return sig_poly, bak_poly, ovr_poly
|
|
|
|
def compute(self, audio, sampling_rate, is_personalized_MOS=False):
|
|
fs = SAMPLING_RATE
|
|
if isinstance(audio, str):
|
|
audio, _ = librosa.load(audio, sr=fs)
|
|
elif sampling_rate != fs:
|
|
|
|
audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=fs)
|
|
actual_audio_len = len(audio)
|
|
len_samples = int(INPUT_LENGTH * fs)
|
|
while len(audio) < len_samples:
|
|
audio = np.append(audio, audio)
|
|
num_hops = int(np.floor(len(audio) / fs) - INPUT_LENGTH) + 1
|
|
hop_len_samples = fs
|
|
predicted_mos_sig_seg_raw = []
|
|
predicted_mos_bak_seg_raw = []
|
|
predicted_mos_ovr_seg_raw = []
|
|
predicted_mos_sig_seg = []
|
|
predicted_mos_bak_seg = []
|
|
predicted_mos_ovr_seg = []
|
|
predicted_p808_mos = []
|
|
|
|
for idx in range(num_hops):
|
|
audio_seg = audio[
|
|
int(idx * hop_len_samples) : int((idx + INPUT_LENGTH) * hop_len_samples)
|
|
]
|
|
if len(audio_seg) < len_samples:
|
|
continue
|
|
input_features = np.array(audio_seg).astype("float32")[np.newaxis, :]
|
|
p808_input_features = np.array(
|
|
self.audio_melspec(audio=audio_seg[:-160])
|
|
).astype("float32")[np.newaxis, :, :]
|
|
oi = {"input_1": input_features}
|
|
p808_oi = {"input_1": p808_input_features}
|
|
p808_mos = self.p808_onnx_sess.run(None, p808_oi)[0][0][0]
|
|
mos_sig_raw, mos_bak_raw, mos_ovr_raw = self.onnx_sess.run(None, oi)[0][0]
|
|
mos_sig, mos_bak, mos_ovr = self.get_polyfit_val(
|
|
mos_sig_raw, mos_bak_raw, mos_ovr_raw, is_personalized_MOS
|
|
)
|
|
predicted_mos_sig_seg_raw.append(mos_sig_raw)
|
|
predicted_mos_bak_seg_raw.append(mos_bak_raw)
|
|
predicted_mos_ovr_seg_raw.append(mos_ovr_raw)
|
|
predicted_mos_sig_seg.append(mos_sig)
|
|
predicted_mos_bak_seg.append(mos_bak)
|
|
predicted_mos_ovr_seg.append(mos_ovr)
|
|
predicted_p808_mos.append(p808_mos)
|
|
clip_dict = {
|
|
"filename": "audio_clip",
|
|
"len_in_sec": actual_audio_len / fs,
|
|
"sr": fs,
|
|
}
|
|
clip_dict["num_hops"] = num_hops
|
|
clip_dict["OVRL_raw"] = np.mean(predicted_mos_ovr_seg_raw)
|
|
clip_dict["SIG_raw"] = np.mean(predicted_mos_sig_seg_raw)
|
|
clip_dict["BAK_raw"] = np.mean(predicted_mos_bak_seg_raw)
|
|
clip_dict["OVRL"] = np.mean(predicted_mos_ovr_seg)
|
|
clip_dict["SIG"] = np.mean(predicted_mos_sig_seg)
|
|
clip_dict["BAK"] = np.mean(predicted_mos_bak_seg)
|
|
clip_dict["P808_MOS"] = np.mean(predicted_p808_mos)
|
|
return clip_dict
|
|
|
