2-D Speech Denoising.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <center>Speech Denoising Using 2-D CNN</center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.7/site-packages/tensorflow_core/python/compat/v2_compat.py:68: disable_resource_variables (from tensorflow.python.ops.variable_scope) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "non-resource variables are not supported in the long term\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf, librosa, numpy as np, pandas as pd, matplotlib.pyplot as plt\n",
    "tf.compat.v1.disable_v2_behavior() # Disabling tensorflow v2.0 behavior\n",
    "np.random.seed(7) # Set random seed for numpy\n",
    "tf.random.set_random_seed(7) # Set tensorflow random seed\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the audio files to spectrograms, matrices of size 513x2459\n",
    "s, sr = librosa.load('train_clean_male.wav', sr = None)\n",
    "S = librosa.stft(s, n_fft = 1024, hop_length = 512)\n",
    "sn, sr = librosa.load('train_dirty_male.wav', sr = None)\n",
    "X = librosa.stft(sn, n_fft = 1024, hop_length = 512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.29152825e-02+0.0000000e+00j  8.45048483e-03+0.0000000e+00j\n",
      "  -2.07481943e-02+0.0000000e+00j ... -1.52135687e-02+0.0000000e+00j\n",
      "  -4.98118671e-03+0.0000000e+00j -7.75823556e-03+0.0000000e+00j]\n",
      " [-7.06597278e-03-4.3706900e-19j -6.94738980e-03+5.1884386e-03j\n",
      "   2.39305142e-02+3.3414264e-03j ...  1.20950164e-02-9.4868029e-03j\n",
      "  -3.12095246e-04-1.4751647e-03j  1.96925160e-02+6.0344348e-03j]\n",
      " [ 3.11613153e-03+2.0498197e-19j -1.59678720e-02-1.7763563e-02j\n",
      "   8.85831192e-03+8.3497362e-03j ... -1.23670641e-02+5.9306147e-03j\n",
      "   7.67358765e-03+2.1299278e-02j -2.82987002e-02+3.5692996e-04j]\n",
      " ...\n",
      " [ 5.43509632e-05+1.1858461e-20j -4.15163022e-03-1.1022503e-03j\n",
      "   1.19022408e-03-2.0613750e-03j ...  5.36997162e-04+7.6431502e-04j\n",
      "  -1.19794244e-02+7.1640010e-03j -9.00672283e-03-8.2842866e-03j]\n",
      " [-4.78351954e-04+4.2351647e-19j  2.63221073e-03-1.5408223e-03j\n",
      "   7.62302428e-04+5.9353854e-03j ... -3.81498365e-03-6.0208812e-03j\n",
      "   5.79493213e-03-3.9956453e-03j  1.48106953e-02+1.0228334e-02j]\n",
      " [ 6.20555074e-04+0.0000000e+00j -3.61816579e-04+0.0000000e+00j\n",
      "  -1.19814684e-03+0.0000000e+00j ...  1.59274472e-03+0.0000000e+00j\n",
      "  -4.55778325e-04+0.0000000e+00j -1.60748102e-02+0.0000000e+00j]]\n"
     ]
    }
   ],
   "source": [
    "# Printing the spectrograms\n",
    "print(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Getting the real valued equivalents of the spectrograms\n",
    "X_mod = abs(X)\n",
    "S_mod = abs(S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(513, 2459) (513, 2459)\n"
     ]
    }
   ],
   "source": [
    "# Dimension of the spectrograms\n",
    "print(X_mod.shape, S_mod.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Getting the transpose matrices for feeding to the 1-D CNN\n",
    "X_T = X_mod.T\n",
    "S_T = S_mod.T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We need to modify our data so that 20 rows of X transpose are the input \"image\" of size 20 x 513 which corresponds to the 20th clean row of S transpose as the \"labels\" for our input image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-7-5864286550b6>:25: DatasetV1.make_initializable_iterator (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `for ... in dataset:` to iterate over a dataset. If using `tf.estimator`, return the `Dataset` object directly from your input function. As a last resort, you can use `tf.compat.v1.data.make_initializable_iterator(dataset)`.\n"
     ]
    }
   ],
   "source": [
    "CNN_2d = tf.Graph() # Graph for the 2-D CNN\n",
    "\n",
    "with CNN_2d.as_default(): \n",
    "\n",
    "    X_dirty = tf.placeholder(tf.float32, shape = [None,513])\n",
    "    S_clean = tf.placeholder(tf.float32, shape = [None,513])\n",
    "    batch_size = tf.placeholder(tf.int64)\n",
    "\n",
    "    # Due to our mapping of the dataet as given above, we need a custom generator function\n",
    "    def gen_data(X_dirty, S_clean):\n",
    "\n",
    "        offset = 0 # To start from\n",
    "        stride = 1 # How much to increment offset by\n",
    "        data_length = X_dirty.shape[0] - 1\n",
    "        while offset <= (data_length-19):\n",
    "            # Create the input image and associated label data\n",
    "            yield(X_dirty[offset:offset+20],S_clean[offset+19])\n",
    "            offset = offset + stride\n",
    "\n",
    "    # Create a tf dataset\n",
    "    dataset = tf.data.Dataset.from_generator(gen_data,(tf.float32,tf.float32),(tf.TensorShape([20,513]),tf.TensorShape([513])),args=(X_dirty,S_clean))\n",
    "    dataset = dataset.repeat()  # Loop through the dataset forever\n",
    "    dataset = dataset.batch(batch_size,drop_remainder=False)  # Batch the newly created data\n",
    "    dataset = dataset.prefetch(1) # Try prefetching\n",
    "    iterator = dataset.make_initializable_iterator()    # Create an iterator for the dataset\n",
    "    next_element = iterator.get_next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-8-9bdb3eaabe63>:20: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.keras.layers.Conv2D` instead.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.7/site-packages/tensorflow_core/python/layers/convolutional.py:424: Layer.apply (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `layer.__call__` method instead.\n",
      "WARNING:tensorflow:From <ipython-input-8-9bdb3eaabe63>:21: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dropout instead.\n",
      "WARNING:tensorflow:From <ipython-input-8-9bdb3eaabe63>:22: max_pooling2d (from tensorflow.python.layers.pooling) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.MaxPooling2D instead.\n",
      "WARNING:tensorflow:From <ipython-input-8-9bdb3eaabe63>:34: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.flatten instead.\n",
      "WARNING:tensorflow:From <ipython-input-8-9bdb3eaabe63>:36: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.Dense instead.\n"
     ]
    }
   ],
   "source": [
    "with CNN_2d.as_default():\n",
    "\n",
    "    (X_nn,y_nn) = iterator.get_next() # Create the 2D CNN network\n",
    "\n",
    "    drop_rate = tf.placeholder(tf.float32)  # Specify the drop rate for the dropout\n",
    "\n",
    "    # Reshape X_nn to have shape = (batch_size,height,width,channels)\n",
    "    X_adj = tf.reshape(X_nn,shape = [batch_size, 20, 513, 1]) # Change batch_size to y_nn.shape[20], last batch has 19 records instead of 20\n",
    "\n",
    "    # Number of filters in each layer and their size\n",
    "    layer1_kernel_num = 15\n",
    "    layer2_kernel_num = 15\n",
    "    layer3_kernel_num = 10\n",
    "\n",
    "    layer1_kernel_size = [5,5]\n",
    "    layer2_kernel_size = [5,5]\n",
    "    layer3_kernel_size = [2,2]\n",
    "\n",
    "    # Complex convolution layer 1\n",
    "    conv1 = tf.layers.conv2d(inputs = X_adj, filters = layer1_kernel_num, kernel_size = layer1_kernel_size, activation = tf.nn.relu, padding='same')\n",
    "    drop1 = tf.layers.dropout(inputs = conv1, rate = drop_rate)\n",
    "    pool1 = tf.layers.max_pooling2d(inputs = drop1, pool_size = [3,3], strides = 2)\n",
    "\n",
    "    # Complex convolution layer 2\n",
    "    conv2 = tf.layers.conv2d(inputs=pool1,filters=layer2_kernel_num,kernel_size=layer2_kernel_size,activation = tf.nn.relu,padding='same')\n",
    "    drop2 = tf.layers.dropout(inputs=conv2,rate=drop_rate)\n",
    "    pool2 = tf.layers.max_pooling2d(inputs=drop2,pool_size=[3,3],strides=2)\n",
    "\n",
    "    # Complex convolution layer 3\n",
    "    conv3 = tf.layers.conv2d(inputs=pool2, filters=layer3_kernel_num, kernel_size=layer3_kernel_size, activation = tf.nn.relu)\n",
    "    drop3 = tf.layers.dropout(inputs=conv3, rate=drop_rate)\n",
    "    pool3 = tf.layers.max_pooling2d(inputs=drop3, pool_size=[2,2], strides=1)\n",
    "\n",
    "    fl = tf.layers.flatten(inputs=pool2)  # Flatten to feed to the fully connected layer\n",
    "\n",
    "    fc1 = tf.layers.dense(inputs=fl, units=800, activation = tf.nn.relu) # Fully connected layer 1\n",
    "\n",
    "    fc = tf.layers.dense(inputs=fc1, units=513, activation = tf.nn.relu)  # Last fully connected layer\n",
    "\n",
    "    loss = tf.reduce_sum(tf.squared_difference(fc, y_nn)) # Define the loss\n",
    "\n",
    "    optimize = tf.train.AdamOptimizer(1e-04).minimize(loss) # Using Adam optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the 2-D CNN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100/100 [==============================] - 908s 9s/step\n"
     ]
    }
   ],
   "source": [
    "with CNN_2d.as_default():\n",
    "   \n",
    "    train_losses = []\n",
    "    sess = tf.Session()\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    sess.run(iterator.initializer, feed_dict = {X_dirty: X_T, S_clean: S_T, batch_size: 20})\n",
    "\n",
    "    max_epochs = 100\n",
    "    progbar = tf.keras.utils.Progbar(max_epochs)\n",
    "    \n",
    "    # Training loop\n",
    "    for ep in range(max_epochs):\n",
    "        batch_losses = 0\n",
    "\n",
    "        for batch in range(122):\n",
    "            loss_, __ = sess.run([loss,optimize], feed_dict = {drop_rate: 0.4, batch_size: 20})\n",
    "            batch_losses = batch_losses + loss_\n",
    "\n",
    "        train_losses.extend([(batch_losses/122)])\n",
    "        progbar.update(ep+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Obtain the predictions for the training audio signal with 2-D CNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2440, 513)\n"
     ]
    }
   ],
   "source": [
    "with CNN_2d.as_default():\n",
    "    sess.run(iterator.initializer, feed_dict = {X_dirty: X_T, S_clean: S_T, batch_size: X_T.shape[0]-19})\n",
    "    X_hat = sess.run(fc, feed_dict = {drop_rate: 0.0, batch_size: X_T.shape[0]-19}) # Get the predicted signal\n",
    "    print(X_hat.shape)  # Get the shape of predicted signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "S_T1_predicted = X_hat.T # Take the transpose\n",
    "S_T1_hat = np.multiply(np.divide(X[:,19:],X_mod[:,19:]),S_T1_predicted) # The predicted S_hat\n",
    "S_T1 = librosa.istft(S_T1_hat,hop_length=512) # Inverse STFT\n",
    "librosa.output.write_wav('train_2D_recon.wav', S_T1, sr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SNR of the training audio signal\n",
    "## SNR = 10 log<sub>10</sub> $\\frac{\\sum_{t}^{} ( s ( t ) )^2} {\\sum_{t}^{} ( s ( t ) - \\hat{s} ( t ) )^2}$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SNR of reconstructed train signal =  20.779985854430667 db\n"
     ]
    }
   ],
   "source": [
    "# Formula to calculate SNR\n",
    "snr = 10 * np.log10(np.sum(np.square(S[:,19:])) / np.sum(np.square(np.subtract(S[:,19:], S_T1_hat))))\n",
    "print(\"SNR of reconstructed train signal = \", abs(snr), \"db\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the 2-D CNN model on test signal 1 (test_x_01.wav)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3, sr = librosa.load('test_x_01.wav', sr = None)\n",
    "T1 = librosa.stft(s3, n_fft=1024, hop_length = 512) # STFT on test_x_01\n",
    "\n",
    "T1_mod = np.abs(T1) # Magnitude of T1\n",
    "T1_T = T1_mod.T     # Transpose\n",
    "with CNN_2d.as_default():\n",
    "    sess.run(iterator.initializer, feed_dict = {X_dirty: T1_T, S_clean: T1_T, batch_size: T1_T.shape[0]-19})\n",
    "    S_T1_predicted_T = sess.run(fc, feed_dict = {drop_rate: 0.0, batch_size: T1_T.shape[0]-19})  # Get the predicted signal\n",
    "S_T1_predicted = S_T1_predicted_T.T  # Take the transpose\n",
    "S_T1_hat = np.multiply(np.divide(T1[:,19:], T1_mod[:,19:]), S_T1_predicted) # The predicted S_hat\n",
    "S_T1 = librosa.istft(S_T1_hat, hop_length = 512)   # Inverse STFT\n",
    "librosa.output.write_wav('test_s_01_2D_CNN_recon.wav', S_T1, sr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the 2-D CNN model on test signal 2 (test_x_02.wav)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3, sr = librosa.load('test_x_02.wav', sr = None)\n",
    "T1 = librosa.stft(s3, n_fft = 1024, hop_length = 512) # STFT on test_x_02\n",
    "\n",
    "T1_mod = np.abs(T1) # Magnitude of T1\n",
    "T1_T = T1_mod.T  # Transpose\n",
    "\n",
    "with CNN_2d.as_default():\n",
    "    sess.run(iterator.initializer, feed_dict = {X_dirty: T1_T, S_clean: T1_T, batch_size: T1_T.shape[0]-19})\n",
    "    S_T1_predicted_T = sess.run(fc, feed_dict = {drop_rate: 0.0, batch_size: T1_T.shape[0]-19})  # Get the predicted signal\n",
    "\n",
    "    S_T1_predicted = S_T1_predicted_T.T  # Transpose\n",
    "S_T1_hat = np.multiply(np.divide(T1[:,19:], T1_mod[:,19:]), S_T1_predicted) # Predicted S_hat\n",
    "S_T1 = librosa.istft(S_T1_hat, hop_length = 512)  # Inverse STFT\n",
    "librosa.output.write_wav('test_s_02_2D_CNN_recon.wav', S_T1, sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.close() # End the session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}