{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Date: Jun 2019<br>Programmer: Christian Dittmar, Yigitcan Özer\n",
    "#### If you use the 'NMF toolbox' please refer to:\n",
    "[1] 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": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import scipy.io.wavfile as wav\n",
    "import matplotlib.pyplot as plt\n",
    "import IPython.display as ipd\n",
    "\n",
    "from NMFtoolbox.forwardSTFT import forwardSTFT\n",
    "from NMFtoolbox.inverseSTFT import inverseSTFT\n",
    "from NMFtoolbox.initTemplates import initTemplates\n",
    "from NMFtoolbox.initActivations import initActivations\n",
    "from NMFtoolbox.NMFD import NMFD\n",
    "from NMFtoolbox.alphaWienerFilter import alphaWienerFilter\n",
    "from NMFtoolbox.visualizeComponentsNMF import visualizeComponentsNMF\n",
    "from NMFtoolbox.utils import make_monaural, pcmInt16ToFloat32Numpy\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",
    "# convert wav from int16 to float32\n",
    "filename = 'runningExample_AmenBreak.wav'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Load the audio signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ozr/.local/lib/python3.6/site-packages/scipy/io/wavfile.py:273: WavFileWarning: Chunk (non-data) not understood, skipping it.\n",
      "  WavFileWarning)\n"
     ]
    }
   ],
   "source": [
    "fs, x = wav.read(os.path.join(inpPath, filename))\n",
    "\n",
    "# make monaural if necessary\n",
    "x = make_monaural(x)\n",
    "\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']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Apply NMF variants to STFT magnitude"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bbecc1a375824bf682c3124de6928bcb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Processing: ', max=30), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# set common parameters\n",
    "numComp = 3\n",
    "numIter = 30\n",
    "numTemplateFrames = 8\n",
    "\n",
    "# generate initial guess for templates\n",
    "paramTemplates = dict()\n",
    "paramTemplates['deltaF'] = deltaF\n",
    "paramTemplates['numComp'] = numComp\n",
    "paramTemplates['numBins'] = numBins\n",
    "paramTemplates['numTemplateFrames'] = numTemplateFrames\n",
    "initW = initTemplates(paramTemplates,'drums')\n",
    "\n",
    "# generate initial activations\n",
    "paramActivations = dict()\n",
    "paramActivations['numComp'] = numComp\n",
    "paramActivations['numFrames'] = numFrames\n",
    "\n",
    "initH = initActivations(paramActivations,'uniform')\n",
    "\n",
    "# NMFD parameters\n",
    "paramNMFD = dict()\n",
    "paramNMFD['numComp'] = numComp\n",
    "paramNMFD['numFrames'] = numFrames\n",
    "paramNMFD['numIter'] = numIter\n",
    "paramNMFD['numTemplateFrames'] = numTemplateFrames\n",
    "paramNMFD['initW'] = initW\n",
    "paramNMFD['initH'] = initH\n",
    "\n",
    "# NMFD core method\n",
    "nmfdW, nmfdH, nmfdV, divKL, _ = NMFD(A, paramNMFD)\n",
    "\n",
    "# alpha-Wiener filtering\n",
    "nmfdA, _ = alphaWienerFilter(A, nmfdV, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 308484 characters]\n",
      "text/plain": [
       "<Figure size 1440x1440 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#visualize\n",
    "paramVis = dict()\n",
    "paramVis['deltaT'] = deltaT\n",
    "paramVis['deltaF'] = deltaF\n",
    "paramVis['endeSec'] = 3.8\n",
    "paramVis['fontSize'] = 24\n",
    "fh1, _ = visualizeComponentsNMF(A, nmfdW, nmfdH, nmfdV, paramVis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "audios = []\n",
    "\n",
    "# resynthesize results of NMF with soft constraints and score information\n",
    "for k in range(numComp):\n",
    "    Y = nmfdA[k] * np.exp(1j * P);\n",
    "    y, _ = inverseSTFT(Y, paramSTFT)\n",
    "    \n",
    "    audios.append(y)\n",
    "    # save result\n",
    "    out_filepath = os.path.join(outPath,\n",
    "                                'Winstons_AmenBreak_NMFD_component_{}.wav'.format(k, filename))\n",
    "    \n",
    "    wav.write(filename=out_filepath, rate=fs, data=y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Kick Drum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 818940 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": [
    "#### Snare Drum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 818940 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": "markdown",
   "metadata": {},
   "source": [
    "#### Ride Cymbal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 818940 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": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(audios[2].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
}
