Add contrast to functional (#551)

* Add contrast to functional

* add tests for contrast and update functional.rst

* Minor changes to sox and batch tests for contrast

docs/source/functional.rst

@@ -93,6 +93,41 @@ Functions to perform common audio operations.
 
 .. autofunction:: equalizer_biquad
 
+:hidden:`bandpass_biquad`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: bandpass_biquad
+
+:hidden:`bandreject_biquad`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: bandreject_biquad
+
+:hidden:`band_biquad`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: band_biquad
+
+:hidden:`treble_biquad`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: treble_biquad
+
+:hidden:`deemph_biquad`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: deemph_biquad
+
+:hidden:`riaa_biquad`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: riaa_biquad
+
+:hidden:`contrast`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: contrast
+
 :hidden:`mask_along_axis`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

test/test_batch_consistency.py

@@ -65,6 +65,10 @@ class TestFunctional(unittest.TestCase):
         ])
         _test_batch(F.istft, stft, n_fft=4, length=4)
 
+    def test_contrast(self):
+        waveform = torch.rand(2, 100) - 0.5
+        _test_batch(F.contrast, waveform, enhancement_amount=80.)
+
 
 class TestTransforms(unittest.TestCase):
     """Test suite for classes defined in `transforms` module"""

test/test_sox_compatibility.py

@@ -300,6 +300,24 @@ class TestFunctionalFiltering(unittest.TestCase):
 
         torch.testing.assert_allclose(output_waveform, sox_output_waveform, atol=1e-4, rtol=1e-5)
 
+    @unittest.skipIf("sox" not in BACKENDS, "sox not available")
+    @AudioBackendScope("sox")
+    def test_contrast(self):
+        """
+        Test contrast effect, compare to SoX implementation
+        """
+        enhancement_amount = 80.
+        noise_filepath = common_utils.get_asset_path('whitenoise.wav')
+        E = torchaudio.sox_effects.SoxEffectsChain()
+        E.set_input_file(noise_filepath)
+        E.append_effect_to_chain("contrast", [enhancement_amount])
+        sox_output_waveform, sr = E.sox_build_flow_effects()
+
+        waveform, sample_rate = torchaudio.load(noise_filepath, normalization=True)
+        output_waveform = F.contrast(waveform, enhancement_amount)
+
+        torch.testing.assert_allclose(output_waveform, sox_output_waveform, atol=1e-4, rtol=1e-5)
+
     @unittest.skipIf("sox" not in BACKENDS, "sox not available")
     @AudioBackendScope("sox")
     def test_equalizer(self):

test/test_torchscript_consistency.py

@@ -406,6 +406,16 @@ class _FunctionalTestMixin:
 
         self._assert_consistency(func, waveform)
 
+    def test_contrast(self):
+        filepath = common_utils.get_asset_path("whitenoise.wav")
+        waveform, _ = torchaudio.load(filepath, normalization=True)
+
+        def func(tensor):
+            enhancement_amount = 80.
+            return F.contrast(tensor, enhancement_amount)
+
+        self._assert_consistency(func, waveform)
+
 
 class _TransformsTestMixin:
     """Implements test for Transforms that are performed for different devices"""

torchaudio/functional.py

@@ -31,6 +31,7 @@ __all__ = [
     "deemph_biquad",
     "riaa_biquad",
     "biquad",
+    "contrast",
     'mask_along_axis',
     'mask_along_axis_iid'
 ]
@@ -1160,6 +1161,38 @@ def riaa_biquad(
     return biquad(waveform, b0, b1, b2, a0, a1, a2)
 
 
+def contrast(
+        waveform: Tensor,
+        enhancement_amount: float = 75.
+) -> Tensor:
+    r"""Apply contrast effect.  Similar to SoX implementation.
+    Comparable with compression, this effect modifies an audio signal to make it sound louder
+
+    Args:
+        waveform (Tensor): audio waveform of dimension of `(..., time)`
+        enhancement_amount (float): controls the amount of the enhancement
+            Allowed range of values for enhancement_amount : 0-100
+            Note that enhancement_amount = 0 still gives a significant contrast enhancement
+
+    Returns:
+        Tensor: Waveform of dimension of `(..., time)`
+
+    References:
+        http://sox.sourceforge.net/sox.html
+    """
+
+    if not 0 <= enhancement_amount <= 100:
+        raise ValueError("Allowed range of values for enhancement_amount : 0-100")
+
+    contrast = enhancement_amount / 750.
+
+    temp1 = waveform * (math.pi / 2)
+    temp2 = contrast * torch.sin(temp1 * 4)
+    output_waveform = torch.sin(temp1 + temp2)
+
+    return output_waveform
+
+
 def mask_along_axis_iid(
         specgrams: Tensor,
         mask_param: int,