Add GPU acceleration (#14)

* added pad clipping flag

* added tensorflow backend

* updated docstrings

* updated readme

* updated readme

* updated package version

* updated readme

* updated readme

* updated website

* updated testing

* updated testing and black formatting

* update testing

* added test requirements

* updated travis yml to install requirements-txt

* updated travis yml

* updated test-req

.travis.yml

@@ -6,4 +6,7 @@ python:
 script:
   - pytest --cov=noisereduce/
 after_success:
-  - coveralls
+  - coveralls
+install:
+  - pip install -r requirements.txt
+  - pip install -r requirements-test.txt

README.md

@@ -23,6 +23,8 @@
 ## Installation
 `pip install noisereduce`
 
+*noisereduce optionally uses Tensorflow as a backend to speed up FFT and gaussian convolution. It is not listed in the requirements.txt so because (1) it is optional and (2) tensorflow-gpu and tensorflow (cpu) are both compatible with this package. The package requires Tensorflow 2+ for all tensorflow operations.* 
+
 ## Usage
 (see notebooks)
 
@@ -45,6 +47,8 @@ win_length (int): Each frame of audio is windowed by `window()`. The window will
 hop_length (int):number audio of frames between STFT columns.
 n_std_thresh (int): how many standard deviations louder than the mean dB of the noise (at each frequency level) to be considered signal
 prop_decrease (float): To what extent should you decrease noise (1 = all, 0 = none)
+pad_clipping (bool): Pad the signals with zeros to ensure that the reconstructed data is equal length to the data
+        use_tensorflow (bool): Use tensorflow as a backend for convolution and fft to speed up computation
 verbose (bool): Whether to plot the steps of the algorithm
 ```
 <div style="text-align:center">

environment.yml

@@ -1,4 +1,4 @@
-name: birdbrain
+name: noisereduce
 channels:
   - conda-forge
   - defaults

noisereduce/noisereduce.py

@@ -3,16 +3,73 @@
 import librosa
 from noisereduce.plotting import plot_reduction_steps
 from tqdm.autonotebook import tqdm
+import warnings
 
+try:
+    import tensorflow as tf
+
+    print(
+        "GPUs available: {}".format(tf.config.experimental.list_physical_devices("GPU"))
+    )
+    if int(tf.__version__[0]) < 2:
+        warnings.warn(
+            "Tensorflow version is below 2.0, some GPU accelerated functionality may not work"
+        )
+except ImportError:
+    warnings.warn(
+        "Tensorflow is not installed and cannot be used for GPU accelerated STFT"
+    )
+
+
+def _stft(y, n_fft, hop_length, win_length, use_tensorflow=False):
+    if use_tensorflow:
+        # return librosa.stft(y=y, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=True)
+        return _stft_tensorflow(y, n_fft, hop_length, win_length)
+    else:
+        return librosa.stft(
+            y=y, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=True
+        )
+
+
+def _istft(y, n_fft, hop_length, win_length, use_tensorflow=False):
+    if use_tensorflow:
+        # return librosa.istft(y, hop_length, win_length)
+        return _istft_tensorflow(y.T, n_fft, hop_length, win_length)
+    else:
+        return librosa.istft(y, hop_length, win_length)
 
-def _stft(y, n_fft, hop_length, win_length):
-    return librosa.stft(y=y, n_fft=n_fft, hop_length=hop_length, win_length=win_length)
+
+def _stft_librosa(y, n_fft, hop_length, win_length):
+    return librosa.stft(
+        y=y, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=True
+    )
 
 
-def _istft(y, hop_length, win_length):
+def _istft_librosa(y, hop_length, win_length):
     return librosa.istft(y, hop_length, win_length)
 
 
+def _stft_tensorflow(y, n_fft, hop_length, win_length):
+    return (
+        tf.signal.stft(
+            y,
+            win_length,
+            hop_length,
+            n_fft,
+            pad_end=True,
+            window_fn=tf.signal.hann_window,
+        )
+        .numpy()
+        .T
+    )
+
+
+def _istft_tensorflow(y, n_fft, hop_length, win_length):
+    return tf.signal.inverse_stft(
+        y.astype(np.complex64), win_length, hop_length, n_fft
+    ).numpy()
+
+
 def _amp_to_db(x):
     return librosa.core.amplitude_to_db(x, ref=1.0, amin=1e-20, top_db=80.0)
 
@@ -31,9 +88,7 @@ def update_pbar(pbar, message):
 
 def _smoothing_filter(n_grad_freq, n_grad_time):
     """Generates a filter to smooth the mask for the spectrogram
-    
-    [description]
-    
+        
     Arguments:
         n_grad_freq {[type]} -- [how many frequency channels to smooth over with the mask.]
         n_grad_time {[type]} -- [how many time channels to smooth over with the mask.]
@@ -58,10 +113,8 @@ def _smoothing_filter(n_grad_freq, n_grad_time):
 
 
 def mask_signal(sig_stft_db, sig_mask, mask_gain_dB, sig_stft):
-    """[summary]
-    
-    [description]
-    
+    """ Reduces amplitude of time/frequency regions of a spectrogram based upon a mask 
+        
     Arguments:
         sig_stft_db {[type]} -- spectrogram of signal in dB
         sig_mask {[type]} -- mask to apply to signal
@@ -83,6 +136,33 @@ def mask_signal(sig_stft_db, sig_mask, mask_gain_dB, sig_stft):
     return sig_stft_amp, sig_stft_db_masked
 
 
+def convolve_gaussian(sig_mask, smoothing_filter, use_tensorflow=False):
+    """ Convolves a gaussian filter with a mask (or any image)
+    
+    Arguments:
+        sig_mask {[type]} -- The signal mask
+        smoothing_filter {[type]} -- the filter to convolve
+    
+    Keyword Arguments:
+        use_tensorflow {bool} -- use tensorflow.signal or scipy.signal (default: {False})
+    """
+    if use_tensorflow:
+        smoothing_filter = smoothing_filter * (
+            (np.shape(smoothing_filter)[1] - 1) / 2 + 1
+        )
+        smoothing_filter = smoothing_filter[:, :, tf.newaxis, tf.newaxis].astype(
+            "float32"
+        )
+        img = sig_mask[:, :, tf.newaxis, tf.newaxis].astype("float32")
+        return (
+            tf.nn.conv2d(img, smoothing_filter, strides=[1, 1, 1, 1], padding="SAME")
+            .numpy()
+            .squeeze()
+        )
+    else:
+        return scipy.signal.fftconvolve(sig_mask, smoothing_filter, mode="same")
+
+
 def reduce_noise(
     audio_clip,
     noise_clip,
@@ -93,6 +173,8 @@ def reduce_noise(
     hop_length=512,
     n_std_thresh=1.5,
     prop_decrease=1.0,
+    pad_clipping=True,
+    use_tensorflow=False,
     verbose=False,
 ):
     """Remove noise from audio based upon a clip containing only noise
@@ -107,6 +189,8 @@ def reduce_noise(
         hop_length (int):number audio of frames between STFT columns.
         n_std_thresh (int): how many standard deviations louder than the mean dB of the noise (at each frequency level) to be considered signal
         prop_decrease (float): To what extent should you decrease noise (1 = all, 0 = none)
+        pad_clipping (bool): Pad the signals with zeros to ensure that the reconstructed data is equal length to the data
+        use_tensorflow (bool): Use tensorflow as a backend for convolution and fft to speed up computation
         verbose (bool): Whether to plot the steps of the algorithm
 
     Returns:
@@ -120,7 +204,9 @@ def reduce_noise(
 
     update_pbar(pbar, "STFT on noise")
     # STFT over noise
-    noise_stft = _stft(noise_clip, n_fft, hop_length, win_length)
+    noise_stft = _stft(
+        noise_clip, n_fft, hop_length, win_length, use_tensorflow=use_tensorflow
+    )
     noise_stft_db = _amp_to_db(np.abs(noise_stft))  # convert to dB
     # Calculate statistics over noise
     update_pbar(pbar, "STFT on signal")
@@ -129,7 +215,15 @@ def reduce_noise(
     noise_thresh = mean_freq_noise + std_freq_noise * n_std_thresh
     # STFT over signal
     update_pbar(pbar, "STFT on signal")
-    sig_stft = _stft(audio_clip, n_fft, hop_length, win_length)
+
+    # pad signal with zeros to avoid extra frames being clipped if desired
+    if pad_clipping:
+        nsamp = len(audio_clip)
+        audio_clip = np.pad(audio_clip, [0, hop_length], mode="constant")
+
+    sig_stft = _stft(
+        audio_clip, n_fft, hop_length, win_length, use_tensorflow=use_tensorflow
+    )
     sig_stft_db = _amp_to_db(np.abs(sig_stft))
     update_pbar(pbar, "Generate mask")
     # Calculate value to mask dB to
@@ -145,19 +239,38 @@ def reduce_noise(
     update_pbar(pbar, "Smooth mask")
     # Create a smoothing filter for the mask in time and frequency
     smoothing_filter = _smoothing_filter(n_grad_freq, n_grad_time)
+
     # convolve the mask with a smoothing filter
+    sig_mask = convolve_gaussian(sig_mask, smoothing_filter, use_tensorflow)
+
     sig_mask = scipy.signal.fftconvolve(sig_mask, smoothing_filter, mode="same")
     sig_mask = sig_mask * prop_decrease
     update_pbar(pbar, "Apply mask")
     # mask the signal
+
     sig_stft_amp, sig_stft_db_masked = mask_signal(
         sig_stft_db, sig_mask, mask_gain_dB, sig_stft
     )
+
     update_pbar(pbar, "Recover signal")
     # recover the signal
-    recovered_signal = _istft(sig_stft_amp, hop_length, win_length)
+    recovered_signal = _istft(
+        sig_stft_amp, n_fft, hop_length, win_length, use_tensorflow=use_tensorflow
+    )
+    # fix the recovered signal length if padding signal
+    if pad_clipping:
+        recovered_signal = librosa.util.fix_length(recovered_signal, nsamp)
+
     recovered_spec = _amp_to_db(
-        np.abs(_stft(recovered_signal, n_fft, hop_length, win_length))
+        np.abs(
+            _stft(
+                recovered_signal,
+                n_fft,
+                hop_length,
+                win_length,
+                use_tensorflow=use_tensorflow,
+            )
+        )
     )
     if verbose:
         plot_reduction_steps(

noisereduce/utils.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+
+def int16_to_float32(data):
+    """ Converts from uint16 wav to float32 wav
+    """
+    if np.max(np.abs(data)) > 32768:
+        raise ValueError("Data has values above 32768")
+    return (data / 32768.0).astype("float32")
+
+
+def float32_to_int16(data):
+    if np.max(data) > 1:
+        data = data / np.max(np.abs(data))
+    return np.array(data * 32767).astype("int16")