declare-lab/tango2

1import numpy as np2 3 4def a_weight(fs, n_fft, min_db=-80.0):5    freq = np.linspace(0, fs // 2, n_fft // 2 + 1)6    freq_sq = np.power(freq, 2)7    freq_sq[0] = 1.08    weight = 2.0 + 20.0 * (2 * np.log10(12194) + 2 * np.log10(freq_sq)9                           - np.log10(freq_sq + 12194 ** 2)10                           - np.log10(freq_sq + 20.6 ** 2)11                           - 0.5 * np.log10(freq_sq + 107.7 ** 2)12                           - 0.5 * np.log10(freq_sq + 737.9 ** 2))13    weight = np.maximum(weight, min_db)14 15    return weight16 17 18def compute_gain(sound, fs, min_db=-80.0, mode="A_weighting"):19    if fs == 16000:20        n_fft = 204821    elif fs == 44100:22        n_fft = 409623    else:24        raise Exception("Invalid fs {}".format(fs))25    stride = n_fft // 226 27    gain = []28    for i in range(0, len(sound) - n_fft + 1, stride):29        if mode == "RMSE":30            g = np.mean(sound[i: i + n_fft] ** 2)31        elif mode == "A_weighting":32            spec = np.fft.rfft(np.hanning(n_fft + 1)[:-1] * sound[i: i + n_fft])33            power_spec = np.abs(spec) ** 234            a_weighted_spec = power_spec * np.power(10, a_weight(fs, n_fft) / 10)35            g = np.sum(a_weighted_spec)36        else:37            raise Exception("Invalid mode {}".format(mode))38        gain.append(g)39 40    gain = np.array(gain)41    gain = np.maximum(gain, np.power(10, min_db / 10))42    gain_db = 10 * np.log10(gain)43    return gain_db44 45 46def mix(sound1, sound2, r, fs):47    gain1 = np.max(compute_gain(sound1, fs))  # Decibel48    gain2 = np.max(compute_gain(sound2, fs))49    t = 1.0 / (1 + np.power(10, (gain1 - gain2) / 20.) * (1 - r) / r)50    sound = ((sound1 * t + sound2 * (1 - t)) / np.sqrt(t ** 2 + (1 - t) ** 2))51    return sound