inoryQwQ's picture
merge AX650 + AX620E(SenseVoice 风格双芯布局)
f722fb4 verified
Raw
History Blame Contribute Delete
19.2 kB
"""RNNoise 信号处理链的 numpy 移植(1:1 对应 origin/rnnoise/src 的 C 实现)。
覆盖 denoise.c / pitch.c / kiss_fft / celt_lpc / rnnoise_tables 的浮点路径:
PCM 帧(480) -> biquad HP -> FFT(960) -> 32 波段能量/DCT + pitch -> 65 维特征
gains/vad -> gain 平滑/限幅 -> pitch filter -> 频谱合成 -> 480 样本
约定:输入音频为 16-bit PCM 等价 float(±32768 量级,不做 /32768 归一化,
与官方 rnnoise demo 一致)。所有运算尽量 float32,与 C float 语义对齐。
"""
from __future__ import annotations
import numpy as np
FRAME_SIZE = 480
WINDOW_SIZE = 2 * FRAME_SIZE
FREQ_SIZE = FRAME_SIZE + 1
NB_BANDS = 32
NB_FEATURES = 2 * NB_BANDS + 1
PITCH_MIN_PERIOD = 60
PITCH_MAX_PERIOD = 768
PITCH_FRAME_SIZE = 960
PITCH_BUF_SIZE = PITCH_MAX_PERIOD + PITCH_FRAME_SIZE
EBAND20MS = np.array(
[0, 2, 4, 6, 8, 10, 12, 15, 18, 21, 24, 28, 32, 36, 41, 47, 53, 60,
68, 77, 87, 98, 110, 124, 140, 157, 176, 198, 223, 251, 282, 317, 356,
400], dtype=np.int32)
SECOND_CHECK = np.array(
[0, 0, 3, 2, 3, 2, 5, 2, 3, 2, 3, 2, 5, 2, 3, 2], dtype=np.int32)
def f32(x):
return np.asarray(x, dtype=np.float32)
def make_half_window() -> np.ndarray:
"""rnn_half_window(dump_rnnoise_tables.c 公式)。"""
i = (np.arange(FRAME_SIZE, dtype=np.float64) + 0.5)
s = np.sin(0.5 * np.pi * i / FRAME_SIZE)
w = np.sin(0.5 * np.pi * s * s)
return f32(w)
HALF_WINDOW = make_half_window()
def make_dct_table() -> np.ndarray:
"""rnn_dct_table:行 i、列 j 为 cos((i+.5)*j*pi/32),j==0 乘 sqrt(.5)。"""
i = np.arange(NB_BANDS, dtype=np.float64)[:, None]
j = np.arange(NB_BANDS, dtype=np.float64)[None, :]
t = np.cos((i + 0.5) * j * np.pi / NB_BANDS)
t[:, 0] *= np.sqrt(0.5)
return f32(t)
DCT_TABLE = make_dct_table()
def fft(x: np.ndarray) -> np.ndarray:
"""forward FFT:kiss_fft 无缩放正变换(960 点),返回复数组。"""
return np.fft.fft(f32(x).astype(np.float64))
def forward_transform(x: np.ndarray) -> np.ndarray:
"""x[960] -> X[481](kiss_fft 该 fork 正变换带 1/N 缩放)。"""
y = fft(x) / WINDOW_SIZE
return y[:FREQ_SIZE]
def inverse_transform(y: np.ndarray) -> np.ndarray:
"""X[481] -> x[960](Hermitian 镜像 + 正变换实现逆变换)。"""
y = np.asarray(y)
n = WINDOW_SIZE
x = np.zeros(n, dtype=np.complex128)
x[:FREQ_SIZE] = y
x[FREQ_SIZE:] = np.conj(y[1:(n - FREQ_SIZE) + 1][::-1])
out = np.fft.fft(x)
res = np.empty(n, dtype=np.float64)
res[0] = out[0].real
res[1:] = out[n:0:-1].real
return res
def apply_window(x: np.ndarray) -> np.ndarray:
x = x.copy()
x[:FRAME_SIZE] *= HALF_WINDOW
x[WINDOW_SIZE - 1:FRAME_SIZE - 1:-1] *= HALF_WINDOW
return x
def compute_band_energy(X: np.ndarray) -> np.ndarray:
"""X[481] 复数 -> bandE[32]。"""
X = np.asarray(X)
power = (X.real ** 2 + X.imag ** 2).astype(np.float64)
# C 实现用 sum[0..NB_BANDS+1],sum[1] 和 sum[NB_BANDS] 修正后取 sum[1..NB_BANDS]
s = np.zeros(NB_BANDS + 2)
for i in range(NB_BANDS):
b0, b1 = int(EBAND20MS[i]), int(EBAND20MS[i + 1])
frac = (np.arange(b1 - b0, dtype=np.float64)) / (b1 - b0)
p = power[b0:b1]
s[i] += np.sum((1 - frac) * p)
s[i + 1] += np.sum(frac * p)
s[1] = (s[0] + s[1]) * 2.0 / 3.0
s[NB_BANDS] = (s[NB_BANDS] + s[NB_BANDS + 1]) * 2.0 / 3.0
return f32(s[1:NB_BANDS + 1])
def compute_band_corr(X: np.ndarray, P: np.ndarray) -> np.ndarray:
X = np.asarray(X)
P = np.asarray(P)
prod = (X.real * P.real + X.imag * P.imag).astype(np.float64)
s = np.zeros(NB_BANDS + 2)
for i in range(NB_BANDS):
b0, b1 = int(EBAND20MS[i]), int(EBAND20MS[i + 1])
frac = (np.arange(b1 - b0, dtype=np.float64)) / (b1 - b0)
p = prod[b0:b1]
s[i] += np.sum((1 - frac) * p)
s[i + 1] += np.sum(frac * p)
s[1] = (s[0] + s[1]) * 2.0 / 3.0
s[NB_BANDS] = (s[NB_BANDS] + s[NB_BANDS + 1]) * 2.0 / 3.0
return f32(s[1:NB_BANDS + 1])
def dct(inp: np.ndarray) -> np.ndarray:
"""out[i] = sqrt(2/22) * sum_j in[j] * table[j, i]。"""
inp = f32(inp)
return f32(inp @ DCT_TABLE * np.sqrt(2.0 / 22.0))
def rnn_biquad(x: np.ndarray, mem: np.ndarray, b, a) -> tuple[np.ndarray, np.ndarray]:
x = f32(x)
y = np.empty_like(x)
mem = mem.copy()
for i in range(x.size):
xi = float(x[i])
yi = xi + mem[0]
mem[0] = mem[1] + (b[0] * xi - a[0] * yi)
mem[1] = b[1] * xi - a[1] * yi
y[i] = yi
return y, mem
def rnn_autocorr(x: np.ndarray, lag: int) -> np.ndarray:
"""rnn_autocorr(overlap=0):ac[k] = sum_i x[i]*x[i+k]。"""
n = x.size
x = f32(x)
ac = np.zeros(lag + 1, dtype=np.float64)
for k in range(lag + 1):
ac[k] = np.dot(x[:n - k].astype(np.float64), x[k:].astype(np.float64))
return ac
def rnn_lpc(ac: np.ndarray, p: int) -> np.ndarray:
ac = np.asarray(ac, dtype=np.float64)
lpc = np.zeros(p, dtype=np.float64)
error = ac[0]
if ac[0] != 0:
for i in range(p):
rr = sum(lpc[j] * ac[i - j] for j in range(i)) + ac[i + 1]
r = -rr / error
lpc[i] = r
for j in range((i + 1) // 2):
tmp1 = lpc[j]
tmp2 = lpc[i - 1 - j]
lpc[j] = tmp1 + r * tmp2
lpc[i - 1 - j] = tmp2 + r * tmp1
error = error - r * r * error
if error < 0.001 * ac[0]:
break
return f32(lpc)
def celt_fir5(x: np.ndarray, num: np.ndarray) -> np.ndarray:
"""5 抽头 FIR:y[i] = x[i] + sum_j num[j]*x[i-1-j]。"""
x = f32(x)
num = f32(num)
n = x.size
y = np.empty(n, dtype=np.float32)
mem = np.zeros(5, dtype=np.float64)
for i in range(n):
xi = float(x[i])
s = xi + float(num[0]) * mem[0] + float(num[1]) * mem[1] + \
float(num[2]) * mem[2] + float(num[3]) * mem[3] + float(num[4]) * mem[4]
mem[4] = mem[3]
mem[3] = mem[2]
mem[2] = mem[1]
mem[1] = mem[0]
mem[0] = xi
y[i] = s
return y
def rnn_pitch_downsample(x: np.ndarray) -> np.ndarray:
"""1728 采样 -> 864 低通(2:1 抽取 + LPC 预加重)。"""
x = f32(x)
length = x.size
half = length // 2
x_lp = np.empty(half, dtype=np.float32)
x_lp[0] = 0.5 * (0.5 * float(x[1]) + float(x[0]))
for i in range(1, half):
x_lp[i] = 0.5 * (0.5 * (float(x[2 * i - 1]) + float(x[2 * i + 1])) + float(x[2 * i]))
ac = rnn_autocorr(x_lp, 4)
ac[0] *= 1.0001
for i in range(1, 5):
ac[i] -= ac[i] * (0.008 * i) * (0.008 * i)
lpc = rnn_lpc(ac, 4)
tmp = 1.0
for i in range(4):
tmp *= 0.9
lpc[i] = lpc[i] * tmp
lpc2 = np.zeros(5, dtype=np.float32)
c1 = 0.8
lpc2[0] = lpc[0] + 0.8
lpc2[1] = lpc[1] + c1 * lpc[0]
lpc2[2] = lpc[2] + c1 * lpc[1]
lpc2[3] = lpc[3] + c1 * lpc[2]
lpc2[4] = c1 * lpc[3]
return celt_fir5(x_lp, lpc2)
def _find_best_pitch(xcorr: np.ndarray, y: np.ndarray, length: int, max_pitch: int):
"""find_best_pitch(浮点路径),返回 best_pitch[2]。"""
xcorr = np.asarray(xcorr, dtype=np.float64)
y = np.asarray(y, dtype=np.float64)
Syy = 1.0 + np.dot(y[:length], y[:length])
best_num = [-1.0, -1.0]
best_den = [0.0, 0.0]
best_pitch = [0, 1]
for i in range(max_pitch):
if xcorr[i] > 0:
xcorr16 = xcorr[i] * 1e-12
num = xcorr16 * xcorr16
if num * best_den[1] > best_num[1] * Syy:
if num * best_den[0] > best_num[0] * Syy:
best_num[1] = best_num[0]
best_den[1] = best_den[0]
best_pitch[1] = best_pitch[0]
best_num[0] = num
best_den[0] = Syy
best_pitch[0] = i
else:
best_num[1] = num
best_den[1] = Syy
best_pitch[1] = i
Syy += y[i + length] * y[i + length] - y[i] * y[i]
Syy = max(1.0, Syy)
return best_pitch
def rnn_pitch_search(x_lp: np.ndarray, y: np.ndarray, length: int, max_pitch: int) -> int:
"""rnn_pitch_search 返回 pitch(粗搜 + 细搜 + 伪插值)。"""
x_lp = f32(x_lp)
y = f32(y)
lag = length + max_pitch
len4 = length >> 2
x_lp4 = x_lp[0:length:2][:len4]
y_lp4 = y[0:lag:2][:(lag >> 2)]
xcorr = np.array([np.dot(x_lp4.astype(np.float64),
y_lp4[i:i + len4].astype(np.float64))
for i in range(max_pitch >> 2)], dtype=np.float64)
best_pitch = _find_best_pitch(xcorr, y_lp4, len4, max_pitch >> 2)
half = max_pitch >> 1
xcorr2 = np.zeros(half, dtype=np.float64)
x_lp64 = x_lp.astype(np.float64)
y64 = y.astype(np.float64)
for i in range(half):
if abs(i - 2 * best_pitch[0]) > 2 and abs(i - 2 * best_pitch[1]) > 2:
continue
s = np.dot(x_lp64[:length >> 1], y64[i:i + (length >> 1)])
xcorr2[i] = max(-1.0, s)
best_pitch = _find_best_pitch(xcorr2, y64, length >> 1, half)
bp = best_pitch[0]
if 0 < bp < half - 1:
a, b, c = xcorr2[bp - 1], xcorr2[bp], xcorr2[bp + 1]
if (c - a) > 0.7 * (b - a):
offset = 1
elif (a - c) > 0.7 * (b - c):
offset = -1
else:
offset = 0
else:
offset = 0
return 2 * bp - offset
def _compute_pitch_gain(xy, xx, yy):
return xy / np.sqrt(1.0 + xx * yy)
def rnn_remove_doubling(x: np.ndarray, maxperiod: int, minperiod: int, n: int,
t0_in: int, prev_period: int, prev_gain: float):
"""rnn_remove_doubling:返回 (pg, T0)。"""
x = f32(x)
minperiod0 = minperiod
maxperiod //= 2
minperiod //= 2
t0 = t0_in // 2
prev_period //= 2
n //= 2
x = x[maxperiod:]
if t0 >= maxperiod:
t0 = maxperiod - 1
T = T0 = t0
x64 = x.astype(np.float64)
xx = np.dot(x64[:n], x64[:n])
xy = np.dot(x64[:n], x64[T0:T0 + n])
yy_lookup = np.zeros(maxperiod + 1, dtype=np.float64)
yy_lookup[0] = xx
yy = xx
for i in range(1, maxperiod + 1):
yy = yy + float(x64[-i]) ** 2 - float(x64[n - i]) ** 2
yy_lookup[i] = max(0.0, yy)
yy = yy_lookup[T0]
best_xy, best_yy = xy, yy
g = g0 = _compute_pitch_gain(xy, xx, yy)
for k in range(2, 16):
T1 = (2 * T0 + k) // (2 * k)
if T1 < minperiod:
break
if k == 2:
T1b = T0 if T1 + T0 > maxperiod else T0 + T1
else:
T1b = (2 * int(SECOND_CHECK[k]) * T0 + k) // (2 * k)
xy2 = np.dot(x64[:n], x64[T1b:T1b + n])
xy = 0.5 * (xy + xy2)
yy = 0.5 * (yy_lookup[T1] + yy_lookup[T1b])
g1 = _compute_pitch_gain(xy, xx, yy)
if abs(T1 - prev_period) <= 1:
cont = prev_gain
elif abs(T1 - prev_period) <= 2 and 5 * k * k < T0:
cont = 0.5 * prev_gain
else:
cont = 0.0
thresh = max(0.3, 0.7 * g0 - cont)
if T1 < 3 * minperiod:
thresh = max(0.4, 0.85 * g0 - cont)
elif T1 < 2 * minperiod:
thresh = max(0.5, 0.9 * g0 - cont)
if g1 > thresh:
best_xy, best_yy = xy, yy
T = T1
g = g1
best_xy = max(0.0, best_xy)
pg = 1.0 if best_yy <= best_xy else best_xy / (best_yy + 1.0)
xcorr = np.zeros(3)
for k in range(3):
lag = T + k - 1
xcorr[k] = np.dot(x64[:n], x64[lag:lag + n])
if (xcorr[2] - xcorr[0]) > 0.7 * (xcorr[1] - xcorr[0]):
offset = 1
elif (xcorr[0] - xcorr[2]) > 0.7 * (xcorr[1] - xcorr[2]):
offset = -1
else:
offset = 0
if pg > g:
pg = g
T0 = 2 * T + offset
if T0 < minperiod0:
T0 = minperiod0
return float(pg), int(T0)
def interp_band_gain(bandE: np.ndarray) -> np.ndarray:
bandE = np.asarray(bandE, dtype=np.float64)
g = np.zeros(FREQ_SIZE)
for i in range(1, NB_BANDS):
b0, b1 = int(EBAND20MS[i]), int(EBAND20MS[i + 1])
frac = (np.arange(b1 - b0, dtype=np.float64)) / (b1 - b0)
g[b0:b1] = (1 - frac) * bandE[i - 1] + frac * bandE[i]
g[:int(EBAND20MS[1])] = bandE[0]
g[int(EBAND20MS[NB_BANDS]):int(EBAND20MS[NB_BANDS + 1])] = bandE[NB_BANDS - 1]
return f32(g)
def rnn_pitch_filter(X, P, Ex, Ep, Exp, g):
"""rnn_pitch_filter:就地修改 X[481]。"""
X = np.asarray(X, dtype=np.complex128)
P = np.asarray(P, dtype=np.complex128)
Ex = np.asarray(Ex, dtype=np.float64)
Ep = np.asarray(Ep, dtype=np.float64)
Exp = np.asarray(Exp, dtype=np.float64)
g = np.asarray(g, dtype=np.float64)
r = np.zeros(NB_BANDS)
for i in range(NB_BANDS):
if Exp[i] > g[i]:
r[i] = 1.0
else:
r[i] = (Exp[i] ** 2) * (1 - g[i] ** 2) / (0.001 + g[i] ** 2 * (1 - Exp[i] ** 2))
r[i] = np.sqrt(min(1.0, max(0.0, r[i])))
r[i] *= np.sqrt(Ex[i] / (1e-8 + Ep[i]))
rf = interp_band_gain(r)
X += rf * P
newE = compute_band_energy(X)
norm = np.sqrt(Ex / (1e-8 + newE))
normf = interp_band_gain(norm)
X *= normf
return X
class RNNoiseState:
"""对应 DenoiseState:维护全部帧间状态。"""
def __init__(self):
self.analysis_mem = np.zeros(FRAME_SIZE, dtype=np.float32)
self.synthesis_mem = np.zeros(FRAME_SIZE, dtype=np.float32)
self.pitch_buf = np.zeros(PITCH_BUF_SIZE, dtype=np.float32)
self.last_gain = 0.0
self.last_period = 0
self.mem_hp_x = np.zeros(2, dtype=np.float32)
self.lastg = np.zeros(NB_BANDS, dtype=np.float32)
self.delayed_X = np.zeros(FREQ_SIZE, dtype=np.complex128)
self.delayed_P = np.zeros(FREQ_SIZE, dtype=np.complex128)
self.delayed_Ex = np.zeros(NB_BANDS, dtype=np.float32)
self.delayed_Ep = np.zeros(NB_BANDS, dtype=np.float32)
self.delayed_Exp = np.zeros(NB_BANDS, dtype=np.float32)
def rnn_frame_analysis(st: RNNoiseState, inp: np.ndarray):
"""返回 (X[481], Ex[32])。"""
x = np.concatenate([st.analysis_mem, f32(inp)])
st.analysis_mem = f32(inp).copy()
x = apply_window(x)
X = forward_transform(x)
Ex = compute_band_energy(X)
return X, Ex
def compute_frame_features(st: RNNoiseState, inp: np.ndarray):
"""返回 (silence, X, P, Ex, Ep, Exp, features)。features 为空时清为 0。"""
X, Ex = rnn_frame_analysis(st, inp)
st.pitch_buf = np.concatenate([st.pitch_buf[FRAME_SIZE:], f32(inp)])
pre = st.pitch_buf
pitch_buf_lp = rnn_pitch_downsample(pre)
x_lp = pitch_buf_lp[(PITCH_MAX_PERIOD >> 1):]
y = pitch_buf_lp
pitch_index = rnn_pitch_search(
x_lp, y, PITCH_FRAME_SIZE, PITCH_MAX_PERIOD - 3 * PITCH_MIN_PERIOD)
pitch_index = PITCH_MAX_PERIOD - pitch_index
gain, pitch_index = rnn_remove_doubling(
st.pitch_buf, PITCH_MAX_PERIOD, PITCH_MIN_PERIOD, PITCH_FRAME_SIZE,
pitch_index, st.last_period, st.last_gain)
st.last_period = pitch_index
st.last_gain = gain
p = np.array([
st.pitch_buf[PITCH_BUF_SIZE - WINDOW_SIZE - pitch_index + i]
for i in range(WINDOW_SIZE)], dtype=np.float32)
p = apply_window(p)
P = forward_transform(p)
Ep = compute_band_energy(P)
Exp = compute_band_corr(X, P)
Exp = Exp / np.sqrt(0.001 + Ex * Ep)
features = np.zeros(NB_FEATURES, dtype=np.float32)
features[NB_BANDS:2 * NB_BANDS] = dct(Exp)
features[2 * NB_BANDS] = 0.01 * (pitch_index - 300)
logMax = -2.0
follow = -2.0
Ly = np.zeros(NB_BANDS)
E = 0.0
for i in range(NB_BANDS):
Ly[i] = np.log10(1e-2 + float(Ex[i]))
Ly[i] = max(logMax - 7, max(follow - 1.5, Ly[i]))
logMax = max(logMax, Ly[i])
follow = max(follow - 1.5, Ly[i])
E += float(Ex[i])
if E < 0.04:
features[:] = 0.0
return 1, X, P, Ex, Ep, Exp, features
features[:NB_BANDS] = dct(Ly)
features[0] -= 12.0
features[1] -= 4.0
return 0, X, P, Ex, Ep, Exp, features
def frame_synthesis(st: RNNoiseState, y) -> np.ndarray:
x = inverse_transform(y)
x = apply_window(x)
out = f32(x[:FRAME_SIZE] + st.synthesis_mem)
st.synthesis_mem = f32(x[FRAME_SIZE:])
return out
def analyze_frame(st: RNNoiseState, inp: np.ndarray) -> dict:
"""C 端 rnnoise_process_frame 前半:biquad + 特征分析。
返回 dict(silence, features, X, P, Ex, Ep, Exp)。会推进 st 的
analysis_mem / pitch_buf / last_period / last_gain 等 DSP 状态。
"""
a_hp = np.array([-1.99599, 0.99600], dtype=np.float32)
b_hp = np.array([-2, 1], dtype=np.float32)
x, st.mem_hp_x = rnn_biquad(f32(inp), st.mem_hp_x, b_hp, a_hp)
silence, X, P, Ex, Ep, Exp, features = compute_frame_features(st, x)
return {"silence": silence, "features": features,
"X": X, "P": P, "Ex": Ex, "Ep": Ep, "Exp": Exp}
def synthesize_frame(st: RNNoiseState, ana: dict, gains, vad):
"""C 端 rnnoise_process_frame 后半:pitch filter + 增益 + 频谱合成。
gains 为 None 表示静音帧(跳过模型相关部分,vad 记 0)。
返回 (out_frame, vad_prob)。需在 analyze_frame 之后调用。
"""
silence = ana["silence"]
X, P, Ex, Ep, Exp = ana["X"], ana["P"], ana["Ex"], ana["Ep"], ana["Exp"]
if not silence and gains is not None:
gains = f32(gains).reshape(-1)
st.delayed_X = rnn_pitch_filter(
st.delayed_X, st.delayed_P, st.delayed_Ex, st.delayed_Ep,
st.delayed_Exp, gains)
for i in range(NB_BANDS):
alpha = 0.6
gains[i] = max(gains[i], alpha * st.lastg[i])
st.lastg[i] = min(1.0, gains[i] * (st.delayed_Ex[i] + 1e-3) / (Ex[i] + 1e-3))
gf = interp_band_gain(gains)
st.delayed_X = st.delayed_X * gf
out = frame_synthesis(st, st.delayed_X)
st.delayed_X = X
st.delayed_P = P
st.delayed_Ex = Ex
st.delayed_Ep = Ep
st.delayed_Exp = Exp
vad_prob = 0.0 if (silence or gains is None) else float(np.asarray(vad).reshape(-1)[0])
return out, vad_prob
def process_frame(st: RNNoiseState, inp: np.ndarray, gains, vad):
"""完整一帧(C 端 rnnoise_process_frame 等价),返回 (out, vad_prob)。"""
ana = analyze_frame(st, inp)
if ana["silence"]:
return synthesize_frame(st, ana, None, 0.0)
return synthesize_frame(st, ana, gains, vad)
def features_from_pcm(pcm: np.ndarray) -> np.ndarray:
"""批量:48k f32 PCM(16-bit 域)-> (T, 65) 特征矩阵(含 silence 清零逻辑)。"""
st = RNNoiseState()
feats = []
pcm = f32(pcm)
for t in range(0, pcm.size - FRAME_SIZE + 1, FRAME_SIZE):
frame = pcm[t:t + FRAME_SIZE]
a_hp = np.array([-1.99599, 0.99600], dtype=np.float32)
b_hp = np.array([-2, 1], dtype=np.float32)
x, st.mem_hp_x = rnn_biquad(frame, st.mem_hp_x, b_hp, a_hp)
_, _, _, _, _, _, features = compute_frame_features(st, x)
feats.append(features.copy())
return np.stack(feats)