File size: 5,835 Bytes
e34c0af |
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 154 155 |
import math
import os
import librosa
import numpy as np
import onnxruntime as ort
from numpy.fft import rfft
from numpy.lib.stride_tricks import as_strided
class PLCMOSEstimator():
def __init__(self, model_version=1):
"""
Initialize a PLC-MOS model of a given version. There are currently three models available, v0 (intrusive)
and v1 (both non-intrusive and intrusive available). The default is to use the v1 models.
"""
self.model_version = model_version
model_paths = [
# v0 model:
[("models/plcmos_v0.onnx", 999999999999), (None, 0)],
# v1 models:
[("models/plcmos_v1_intrusive.onnx", 768),
("models/plcmos_v1_nonintrusive.onnx", 999999999999)],
]
self.sessions = []
self.max_lens = []
options = ort.SessionOptions()
options.intra_op_num_threads = 8
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
for path, max_len in model_paths[model_version]:
if not path is None:
file_dir = os.path.dirname(os.path.realpath(__file__))
self.sessions.append(ort.InferenceSession(
os.path.join(file_dir, path), options))
self.max_lens.append(max_len)
else:
self.sessions.append(None)
self.max_lens.append(0)
def logpow_dns(self, sig, floor=-30.):
"""
Compute log power of complex spectrum.
Floor any -`np.inf` value to (nonzero minimum + `floor`) dB.
If all values are 0s, floor all values to -80 dB.
"""
log10e = np.log10(np.e)
pspec = sig.real ** 2 + sig.imag ** 2
zeros = pspec == 0
logp = np.empty_like(pspec)
if np.any(~zeros):
logp[~zeros] = np.log(pspec[~zeros])
logp[zeros] = np.log(pspec[~zeros].min()) + floor / 10 / log10e
else:
logp.fill(-80 / 10 / log10e)
return logp
def hop2hsize(self, wind, hop):
"""
Convert hop fraction to integer size if necessary.
"""
if hop >= 1:
assert type(hop) == int, "Hop size must be integer!"
return hop
else:
assert 0 < hop < 1, "Hop fraction has to be in range (0,1)!"
return int(len(wind) * hop)
def stana(self, sig, sr, wind, hop, synth=False, center=False):
"""
Short term analysis by windowing
"""
ssize = len(sig)
fsize = len(wind)
hsize = self.hop2hsize(wind, hop)
if synth:
sstart = hsize - fsize # int(-fsize * (1-hfrac))
elif center:
sstart = -int(len(wind) / 2) # odd window centered at exactly n=0
else:
sstart = 0
send = ssize
nframe = math.ceil((send - sstart) / hsize)
# Calculate zero-padding sizes
zpleft = -sstart
zpright = (nframe - 1) * hsize + fsize - zpleft - ssize
if zpleft > 0 or zpright > 0:
sigpad = np.zeros(ssize + zpleft + zpright, dtype=sig.dtype)
sigpad[zpleft:len(sigpad) - zpright] = sig
else:
sigpad = sig
return as_strided(sigpad, shape=(nframe, fsize),
strides=(sig.itemsize * hsize, sig.itemsize)) * wind
def stft(self, sig, sr, wind, hop, nfft):
"""
Compute STFT: window + rfft
"""
frames = self.stana(sig, sr, wind, hop, synth=True)
return rfft(frames, n=nfft)
def stft_transform(self, audio, dft_size=512, hop_fraction=0.5, sr=16000):
"""
Compute STFT parameters, then compute STFT
"""
window = np.hamming(dft_size + 1)
window = window[:-1]
amp = np.abs(self.stft(audio, sr, window, hop_fraction, dft_size))
feat = self.logpow_dns(amp, floor=-120.)
return feat / 20.
def run(self, audio_degraded, audio_clean=None, combined=False):
"""
Run the PLCMOS model and return the MOS for the given audio. If a clean audio file is passed and the
selected model version has an intrusive version, that version will be used, otherwise, the nonintrusive
model will be used. If combined is set to true (default), the mean of intrusive and nonintrusive models
results will be returned, when both are available
For intrusive models, the clean reference should be the unprocessed audio file the degraded audio is
based on. It is not required to be aligned with the degraded audio.
Audio data should be 16kHz, mono, [-1, 1] range.
"""
audio_features_degraded = np.float32(self.stft_transform(audio_degraded))[
np.newaxis, np.newaxis, ...]
assert len(
audio_features_degraded) <= self.max_lens[0], "Maximum input length exceeded"
if audio_clean is None:
combined = False
mos = 0
session = self.sessions[0]
assert not session is None, "Intrusive model not available for this model version."
audio_features_clean = np.float32(self.stft_transform(audio_clean))[
np.newaxis, np.newaxis, ...]
assert len(
audio_features_clean) <= self.max_lens[0], "Maximum input length exceeded"
onnx_inputs = {"degraded_audio": audio_features_degraded,
"clean_audio": audio_features_clean}
mos = float(session.run(None, onnx_inputs)[0])
session = self.sessions[1]
assert not session is None, "Nonintrusive model not available for this model version."
onnx_inputs = {"degraded_audio": audio_features_degraded}
mos_2 = float(session.run(None, onnx_inputs)[0])
mos = [mos, mos_2]
return mos
|