File size: 4,173 Bytes
c760209 | 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 | import numpy as np
from math import gcd
SR = 48000
N_FFT = 2048
HOP = 480
N_MELS = 256
FMIN = 20.0
FMAX = 24000.0
LOG_FLOOR = 1e-5
MAG_EPS = 1e-9
_MEL = None
_WIN = None
def _hz_to_mel(f):
f = np.asarray(f, dtype=np.float64)
f_sp, min_log_hz = 200.0 / 3.0, 1000.0
min_log_mel = min_log_hz / f_sp
logstep = np.log(6.4) / 27.0
mel = np.empty_like(f)
lin = f < min_log_hz
mel[lin] = f[lin] / f_sp
mel[~lin] = min_log_mel + np.log(f[~lin] / min_log_hz) / logstep
return mel
def _mel_to_hz(m):
m = np.asarray(m, dtype=np.float64)
f_sp, min_log_hz = 200.0 / 3.0, 1000.0
min_log_mel = min_log_hz / f_sp
logstep = np.log(6.4) / 27.0
hz = np.empty_like(m)
lin = m < min_log_mel
hz[lin] = m[lin] * f_sp
hz[~lin] = min_log_hz * np.exp(logstep * (m[~lin] - min_log_mel))
return hz
def mel_filterbank():
global _MEL
if _MEL is None:
freqs = np.linspace(0.0, SR / 2.0, N_FFT // 2 + 1)
edges = np.linspace(_hz_to_mel([FMIN])[0], _hz_to_mel([FMAX])[0], N_MELS + 2)
hz = _mel_to_hz(edges)
fb = np.zeros((N_MELS, freqs.shape[0]), dtype=np.float32)
for m in range(N_MELS):
left, center, right = hz[m], hz[m + 1], hz[m + 2]
up = (freqs - left) / (center - left)
down = (right - freqs) / (right - center)
fb[m] = np.maximum(0.0, np.minimum(up, down)) * (2.0 / max(1e-8, right - left))
_MEL = fb
return _MEL
def hann():
global _WIN
if _WIN is None:
_WIN = np.hanning(N_FFT + 1)[:-1].astype(np.float64)
return _WIN
def log_mel(audio):
x = np.asarray(audio, dtype=np.float64).reshape(-1)
pad = (N_FFT - HOP) // 2
x = np.pad(x, (pad, pad), mode="reflect")
frames = 1 + (x.size - N_FFT) // HOP
idx = np.arange(N_FFT)[None, :] + HOP * np.arange(frames)[:, None]
spec = np.fft.rfft(x[idx] * hann(), n=N_FFT, axis=1).T
mag = np.sqrt(spec.real ** 2 + spec.imag ** 2 + MAG_EPS)
mel = np.log(np.clip(mel_filterbank() @ mag, LOG_FLOOR, None))
return mel.T[None].astype(np.float32)
def stft(x, center=True, pad_mode="constant"):
x = np.asarray(x, dtype=np.float64).reshape(-1)
if center:
p = N_FFT // 2
x = np.pad(x, (p, p), mode=pad_mode)
frames = 1 + (x.size - N_FFT) // HOP
idx = np.arange(N_FFT)[None, :] + HOP * np.arange(frames)[:, None]
return np.fft.rfft(x[idx] * hann(), n=N_FFT, axis=1).T
def istft(spec, length=None, center=True):
spec = np.asarray(spec)
frames = np.fft.irfft(spec.T, n=N_FFT, axis=1) * hann()
out_len = N_FFT + HOP * (frames.shape[0] - 1)
ola = np.zeros(out_len, dtype=np.float64)
wsum = np.zeros(out_len, dtype=np.float64)
w2 = hann() ** 2
for i in range(frames.shape[0]):
s = i * HOP
ola[s:s + N_FFT] += frames[i]
wsum[s:s + N_FFT] += w2
ola /= np.where(wsum > 1e-11, wsum, 1.0)
if center:
ola = ola[N_FFT // 2:len(ola) - N_FFT // 2]
if length is not None:
if ola.size > length:
ola = ola[:length]
elif ola.size < length:
ola = np.pad(ola, (0, length - ola.size))
return ola.astype(np.float32)
def cutoff_bin(spec):
energy = np.cumsum(np.abs(spec).sum(axis=-1))
if energy[-1] <= 0:
return 0
limit = energy[-1] * 0.99
below = np.nonzero(energy < limit)[0]
return int(below[-1]) + 1 if below.size else 0
def postprocess(pred, src, length):
ps = stft(pred, center=True, pad_mode="constant")
ss = stft(src, center=True, pad_mode="constant")
n = min(ps.shape[-1], ss.shape[-1])
ps, ss = ps[:, :n], ss[:, :n]
cr = cutoff_bin(ss)
merged = ps.copy()
merged[:cr] = ss[:cr]
out = istft(merged, length=length, center=True)
peak = np.max(np.abs(out))
if peak > 0:
out = out / peak * 0.99
return out.astype(np.float32)
def resample(x, sr_in, sr_out=SR):
x = np.asarray(x, dtype=np.float32).reshape(-1)
if sr_in == sr_out:
return x
g = gcd(int(sr_out), int(sr_in))
from scipy.signal import resample_poly
return resample_poly(x, sr_out // g, sr_in // g).astype(np.float32)
|