diff --git a/test/test_functional_filtering.py b/test/test_functional_filtering.py index b3c23a67b7..ab209ed7a7 100644 --- a/test/test_functional_filtering.py +++ b/test/test_functional_filtering.py @@ -136,6 +136,26 @@ def test_highpass(self): # TBD - this fails at the 1e-4 level, debug why assert torch.allclose(sox_output_waveform, output_waveform, atol=1e-3) + def test_equalizer(self): + """ + Test biquad peaking equalizer filter, compare to SoX implementation + """ + + CENTER_FREQ = 300 + Q = 0.707 + GAIN = 1 + + noise_filepath = os.path.join(self.test_dirpath, "assets", "whitenoise.mp3") + E = torchaudio.sox_effects.SoxEffectsChain() + E.set_input_file(noise_filepath) + E.append_effect_to_chain("equalizer", [CENTER_FREQ, Q, GAIN]) + sox_output_waveform, sr = E.sox_build_flow_effects() + + waveform, sample_rate = torchaudio.load(noise_filepath, normalization=True) + output_waveform = F.equalizer_biquad(waveform, sample_rate, CENTER_FREQ, GAIN, Q) + + assert torch.allclose(sox_output_waveform, output_waveform, atol=1e-4) + def test_perf_biquad_filtering(self): fn_sine = os.path.join(self.test_dirpath, "assets", "whitenoise.mp3") diff --git a/torchaudio/functional.py b/torchaudio/functional.py index 6fa4b5341e..92d1ac2410 100644 --- a/torchaudio/functional.py +++ b/torchaudio/functional.py @@ -17,6 +17,7 @@ "lfilter", "lowpass_biquad", "highpass_biquad", + "equalizer_biquad", "biquad", 'mask_along_axis', 'mask_along_axis_iid' @@ -683,6 +684,33 @@ def lowpass_biquad(waveform, sample_rate, cutoff_freq, Q=0.707): return biquad(waveform, b0, b1, b2, a0, a1, a2) +def equalizer_biquad(waveform, sample_rate, center_freq, gain, Q=0.707): + # type: (Tensor, int, float, float, float) -> Tensor + r"""Designs biquad peaking equalizer filter and performs filtering. Similar to SoX implementation. + + Args: + waveform (torch.Tensor): audio waveform of dimension of `(channel, time)` + sample_rate (int): sampling rate of the waveform, e.g. 44100 (Hz) + center_freq (float): filter’s central frequency + gain (float): desired gain at the boost (or attenuation) in dB + q_factor (float): https://en.wikipedia.org/wiki/Q_factor + + Returns: + output_waveform (torch.Tensor): Dimension of `(channel, time)` + """ + w0 = 2 * math.pi * center_freq / sample_rate + A = math.exp(gain / 40.0 * math.log(10)) + alpha = math.sin(w0) / 2 / Q + + b0 = 1 + alpha * A + b1 = -2 * math.cos(w0) + b2 = 1 - alpha * A + a0 = 1 + alpha / A + a1 = -2 * math.cos(w0) + a2 = 1 - alpha / A + return biquad(waveform, b0, b1, b2, a0, a1, a2) + + @torch.jit.script def mask_along_axis_iid(specgrams, mask_param, mask_value, axis): # type: (Tensor, int, float, int) -> Tensor