{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Date: Jun 2019<br>Programmer: Christian Dittmar, Yiğitcan Özer\n",
    "This is the demo script which illustrates the main functionalities of the 'NMF toolbox'. For a detailed description we refer to [1,2] (see References below).\n",
    "\n",
    "#### The script proceeds in the following steps:\n",
    "<br>1. It loads an example audio file containing drums and melodic instruments\n",
    "<br>2. It computes the STFT of the audio data.\n",
    "<br>3. It applies KAM and NMF as described in [2], with random initialization of the NMF components. The final classification into harmonic and percussive is done according to the percussiveness threshold p_thresh = 0.25 as given in [2].\n",
    "<br>4. It visualizes the decomposition results.\n",
    "<br>5. It resynthesizes the separated audio streams and saves them as wav files to the hard drive.\n",
    "\n",
    "#### Reference: \n",
    "[1] Christian Dittmar, Meinard Müller\n",
    "<br>Reverse Engineering the Amen Break - Score-informed Separation and\n",
    "<br>Restoration applied to Drum Recordings\n",
    "<br>IEEE/ACM Transactions on Audio, Speech, and Language Processing,\n",
    "<br>24(9): 1531-1543, 2016.\n",
    "<br>[2] Christian Dittmar, Patricio López-Serrano, Meinard Müller<br>Unifying Local and Global Methods for Harmonic-Percussive Source Separation<br>In Proceedings of the IEEE International Conference on Acoustics,<br>Speech, and Signal Processing (ICASSP), 2018.\n",
    "\n",
    "#### If you use the 'NMF toolbox' please refer to:\n",
    "[3] Patricio López-Serrano, Christian Dittmar, Yiğitcan Özer, and Meinard Müller\n",
    "<br>NMF Toolbox: Music Processing Applications of Nonnegative Matrix Factorization\n",
    "<br>In Proceedings of the International Conference on Digital Audio Effects (DAFx), 2019.\n",
    "\n",
    "#### License:\n",
    "This file is part of 'NMF toolbox'. 'NMF toolbox' is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. 'NMF toolbox' is distributed in the hope that it will be useful, but ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.\n",
    "\n",
    "#### You should have received a copy of the GNU General Public License along<br>with 'NMF toolbox'. If not, see http://www.gnu.org/licenses/."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import scipy.io.wavfile as wav\n",
    "import IPython.display as ipd\n",
    "\n",
    "from NMFtoolbox.utils import make_monaural, pcmInt16ToFloat32Numpy\n",
    "from NMFtoolbox.forwardSTFT import forwardSTFT\n",
    "from NMFtoolbox.inverseSTFT import inverseSTFT\n",
    "from NMFtoolbox.logFreqLogMag import logFreqLogMag\n",
    "from NMFtoolbox.HPSS_KAM import HPSS_KAM_Fitzgerald\n",
    "from NMFtoolbox.initTemplates import initTemplates\n",
    "from NMFtoolbox.initActivations import initActivations\n",
    "from NMFtoolbox.drumSpecificSoftConstraintsNMF import drumSpecificSoftConstraintsNMF\n",
    "from NMFtoolbox.NMFD import NMFD\n",
    "from NMFtoolbox.percussivenessEstimation import percussivenessEstimation\n",
    "from NMFtoolbox.convModel import convModel\n",
    "from NMFtoolbox.alphaWienerFilter import alphaWienerFilter\n",
    "from NMFtoolbox.visualizeComponentsNMF import visualizeComponentsNMF\n",
    "from NMFtoolbox.visualizeComponentsKAM import visualizeComponentsKAM\n",
    "\n",
    "inpPath = '../data/'\n",
    "outPath = 'output/'\n",
    "\n",
    "# create the output directory if it doesn't exist\n",
    "if not os.path.isdir(outPath):\n",
    "    os.makedirs(outPath)\n",
    "\n",
    "filename = 'runningExample_IGotYouMixture.wav'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Load the audio signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read signal\n",
    "fs, x = wav.read(os.path.join(inpPath, filename))\n",
    "\n",
    "# make monaural if necessary\n",
    "x = make_monaural(x)\n",
    "\n",
    "# convert wav from int16 to float32\n",
    "x = pcmInt16ToFloat32Numpy(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Compute STFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# spectral parameters\n",
    "paramSTFT = dict()\n",
    "paramSTFT['blockSize'] = 2048\n",
    "paramSTFT['hopSize'] = 512\n",
    "paramSTFT['winFunc'] = np.hanning(paramSTFT['blockSize'])\n",
    "paramSTFT['reconstMirror'] = True\n",
    "paramSTFT['appendFrame'] = True\n",
    "paramSTFT['numSamples'] = len(x)\n",
    "\n",
    "# STFT computation\n",
    "X, A, P = forwardSTFT(x, paramSTFT)\n",
    "\n",
    "# get dimensions and time and freq resolutions\n",
    "numBins, numFrames = X.shape\n",
    "deltaT = paramSTFT['hopSize'] / fs\n",
    "deltaF = fs / paramSTFT['blockSize']\n",
    "\n",
    "# get logarithmically-spaced frequency axis version for visualization purposes\n",
    "logFreqLogMagA, logFreqAxis = logFreqLogMag(A, deltaF)\n",
    "numLogBins = len(logFreqAxis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Apply KAM-based Harmonic Percussive Separation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 488080 characters]\n",
      "text/plain": [
       "<Figure size 1080x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# set common parameters\n",
    "numIterKAM = 30\n",
    "kamA, Kern, KernOrd = HPSS_KAM_Fitzgerald(A, numIterKAM, 13)\n",
    "\n",
    "# visualize\n",
    "paramVis = dict()\n",
    "paramVis['deltaT'] = deltaT\n",
    "paramVis['deltaF'] = deltaF\n",
    "paramVis['fontSize'] = 14\n",
    "\n",
    "fh1 = visualizeComponentsKAM(kamA, paramVis)\n",
    "\n",
    "# save result\n",
    "fh1.savefig(os.path.join(outPath, 'demoDrumExtractionKAM_NMF_percThreshold_KAM.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "audios = []\n",
    "\n",
    "# resynthesize KAM results\n",
    "for k in range(2):\n",
    "    Y = kamA[k] * np.exp(1j * P);\n",
    "    y, _ = inverseSTFT(Y, paramSTFT)\n",
    "    audios.append(y)\n",
    "    # save result\n",
    "    out_filepath = os.path.join(outPath,\n",
    "                                'demoDrumExtractionKAM_NMF_percThreshold_KAM_component_{}_extracted_from_{}'.format(k, filename))\n",
    "    \n",
    "    wav.write(filename=out_filepath, rate=fs, data=y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Input audio mixture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 792684 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(x, rate=fs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### KAM-based percussive component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 792684 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(audios[0].T, rate=fs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### KAM-based harmonic component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 792684 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(audios[1].T, rate=fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# concatenate new NMF target\n",
    "V = np.concatenate([kamA[0], kamA[1]])\n",
    "numDoubleBins = V.shape[0]\n",
    "\n",
    "# prepare matrix to revert concatenation,\n",
    "AccuMat = np.concatenate([np.eye(numBins), np.eye(numBins)], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Apply NMF with drum-specific soft constraints to KAM-based target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6a1f8f12b9d4847b184426850fab2d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Processing', max=60, style=ProgressStyle(description_width='i…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# set common parameters\n",
    "numIterNMF = 60\n",
    "numComp = 30\n",
    "numTemplateFrames = 1\n",
    "\n",
    "# generate random templates covering 2 times the original frequency range\n",
    "paramTemplates = dict()\n",
    "paramTemplates['numComp'] = numComp\n",
    "paramTemplates['numBins'] = numDoubleBins\n",
    "paramTemplates['numTemplateFrames'] = numTemplateFrames\n",
    "initW = initTemplates(paramTemplates, 'random')\n",
    "\n",
    "# generate uniform activations\n",
    "paramActivations = dict()\n",
    "paramActivations['numComp'] = numComp\n",
    "paramActivations['numFrames'] = numFrames\n",
    "initH = initActivations(paramActivations, 'uniform')\n",
    "\n",
    "# set NMFD parameters\n",
    "paramNMFD = dict()\n",
    "paramNMFD['numComp'] = numComp\n",
    "paramNMFD['numFrames'] = numFrames\n",
    "paramNMFD['numIter'] = numIterNMF\n",
    "paramNMFD['numTemplateFrames'] = numTemplateFrames\n",
    "paramNMFD['numBins'] = numDoubleBins\n",
    "paramNMFD['initW'] = initW\n",
    "paramNMFD['initH'] = initH\n",
    "\n",
    "# set soft constraint parameters\n",
    "paramConstr = dict()\n",
    "paramConstr['funcPointerPreProcess'] = drumSpecificSoftConstraintsNMF\n",
    "paramConstr['Kern'] = Kern\n",
    "paramConstr['KernOrd'] = KernOrd\n",
    "paramConstr['decay'] = 0.75\n",
    "paramConstr['numBinsDrum'] = numBins\n",
    "\n",
    "# NMFD core method\n",
    "nmfdW, nmfdH, _, _, tensorW = NMFD(V, paramNMFD, paramConstr)\n",
    "\n",
    "# get final percussiveness estimate\n",
    "percWeight = percussivenessEstimation(tensorW)\n",
    "\n",
    "# re-order components, only for visualization\n",
    "idx_sorted_rev = np.argsort(-percWeight)\n",
    "\n",
    "tensorW = tensorW[:, idx_sorted_rev, :]\n",
    "nmfdH = nmfdH[idx_sorted_rev, :]\n",
    "nmfdW = [nmfdW[idx] for idx in idx_sorted_rev]\n",
    "percWeight = percWeight[idx_sorted_rev]\n",
    "\n",
    "# perform final thresholding\n",
    "percWeight = np.where(percWeight > 0.25, 1.0, 0.0)\n",
    "\n",
    "# compute separate models for percussive and harmonic part\n",
    "# in the case of numTemplateFrames=1, this step equals eq. (1) in [2]\n",
    "Vp = convModel(tensorW, np.diag(percWeight) @ nmfdH)\n",
    "Vh = convModel(tensorW, np.diag(1 - percWeight) @ nmfdH)\n",
    "\n",
    "# accumulate back to original spectrum, reverting the stacking\n",
    "# this step is described in the last paragraph of sec. 2.4 in [2]\n",
    "Ap = AccuMat @ Vp\n",
    "Ah = AccuMat @ Vh\n",
    "\n",
    "# alpha-Wiener filtering\n",
    "nmfdA, _ = alphaWienerFilter(A, [Ap, Ah], 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 571740 characters]\n",
      "text/plain": [
       "<Figure size 1440x1440 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# create reduced version of templates for visualization\n",
    "paramVis['fontSize'] = 24\n",
    "\n",
    "nmfdW_vis = list()\n",
    "for nmfdW_curr in nmfdW:\n",
    "    nmfdW_curr = AccuMat @ nmfdW_curr\n",
    "    nmfdW_vis.append(nmfdW_curr)\n",
    "\n",
    "# nmfdW = AccuMat @ np.array(nmfdW).squeeze().T\n",
    "fh2, _ = visualizeComponentsNMF(A, nmfdW_vis, nmfdH, nmfdA, paramVis);\n",
    "\n",
    "# save result\n",
    "fh2.savefig(os.path.join(outPath, 'demoDrumExtractionKAM_NMF_percThreshold_NMF.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "audios = []\n",
    "\n",
    "# resynthesize NMF with soft constraints results\n",
    "for k in range(2):\n",
    "    Y = nmfdA[k] * np.exp(1j * P);\n",
    "    y, _ = inverseSTFT(Y, paramSTFT)\n",
    "    audios.append(y)\n",
    "    \n",
    "    # save result\n",
    "    out_filepath = os.path.join(outPath,\n",
    "                                'demoDrumExtractionKAM_NMF_percThreshold_NMF_component_{}_extracted_from_{}'.format(k, filename))\n",
    "    \n",
    "    wav.write(filename=out_filepath, rate=fs, data=y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Input audio mixture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 792684 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(x, rate=fs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Percussive component based on KAM + NMF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 792684 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(audios[0].T, rate=fs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Harmonic component based on KAM + NMF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 792684 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(audios[1].T, rate=fs)"
   ]
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
