audiomodel_final_clear.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 122
    },
    "colab_type": "code",
    "id": "HIsMREXhm2Jf",
    "outputId": "8a3d1556-3e95-427a-95b2-3d93b4ef626b"
   },
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "afwl9L_zLjrs"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "import os, sys\n",
    "from os.path import exists\n",
    "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
    "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
    "cuda_output = !ldconfig -p|grep cudart.so|sed -e 's/.*\\.\\([0-9]*\\)\\.\\([0-9]*\\)$/cu\\1\\2/'\n",
    "accelerator = cuda_output[0] if exists('/dev/nvidia0') else 'cpu'\n",
    "!pip install torch torchaudio torchvision\n",
    "!pip install livelossplot\n",
    "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
    "!pip install -U nvgpu\n",
    "\n",
    "import torch\n",
    "import torchaudio\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from livelossplot import PlotLosses\n",
    "import math\n",
    "import numpy as np\n",
    "import torchvision\n",
    "from IPython.display import clear_output\n",
    "from IPython.display import Audio\n",
    "import librosa\n",
    "import os\n",
    "import pynvml\n",
    "\n",
    "torch.set_printoptions(precision=7)\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xDY3jhBDPPk3"
   },
   "outputs": [],
   "source": [
    "# a function to track the the available GPU memory\n",
    "def memory():\n",
    "  pynvml.nvmlInit()\n",
    "  GPU_ID = 0\n",
    "  handle = pynvml.nvmlDeviceGetHandleByIndex(GPU_ID)\n",
    "  meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)\n",
    "  MB_SIZE = 1024*1024\n",
    "  print('available: ',round(meminfo.total/MB_SIZE)) # 6078 MB\n",
    "  print('used:      ',round(meminfo.used/MB_SIZE))  # 531 MB\n",
    "  print('free:      ',round(meminfo.free/MB_SIZE))  # 5546 MB\n",
    "\n",
    "  pynvml.nvmlShutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UlcwRln_PRCN"
   },
   "outputs": [],
   "source": [
    "# load the normalized audio clips into tensors\n",
    "\n",
    "bongo_dir = \"drive/My Drive/Colab Notebooks/Bachata/1 bar/bongo/normalized\"\n",
    "guira_dir = \"drive/My Drive/Colab Notebooks/Bachata/1 bar/guira/normalized\"\n",
    "bass_dir = \"drive/My Drive/Colab Notebooks/Bachata/1 bar/bass/normalized\"\n",
    "guitar_dir = \"drive/My Drive/Colab Notebooks/Bachata/1 bar/guitar/normalized\"\n",
    "directories = [bongo_dir,guira_dir,bass_dir,guitar_dir]\n",
    "instruments = ['bongo','guira','bass','guitar']\n",
    "  \n",
    "\n",
    "total_list=[]\n",
    "for directory in directories:\n",
    "  mylist = []\n",
    "  for filename in os.listdir(directory):\n",
    "    if filename.endswith(\".aiff\"):\n",
    "      myfile = torchaudio.load(\"{0}/{1}\".format(directory,filename))[0][0]\n",
    "      if len(myfile) == 36000:          #splitting the 1 bar loops into 2 half bar loops\n",
    "        mylist.append(myfile[:18000])\n",
    "        mylist.append(myfile[18000:])\n",
    "  total_list.append(mylist)\n",
    "bongo_list = torch.stack(total_list[0])\n",
    "guira_list = torch.stack(total_list[1])\n",
    "bass_list = torch.stack(total_list[2])\n",
    "guitar_list = torch.stack(total_list[3])\n",
    "instruments_wave = [bongo_list,guira_list,bass_list,guitar_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YC9Xe_l7kwWN"
   },
   "outputs": [],
   "source": [
    "# obtaining the magnitude and phase matrices via short-time Fourier transform\n",
    "def _stft(waveform):\n",
    "    waveform = waveform.squeeze(0)\n",
    "    spec = torch.stft(waveform,n_fft=510,hop_length=70,win_length=510,window=torch.hann_window(510))  #short time Fourier transform to produce a complex matrix\n",
    "    mag = (spec[...,0] * spec[...,0] + spec[...,1] * spec[...,1])**0.5 # piecewise modulus of the complex stft matrix to obtain the amplitude/magnitude matrix\n",
    "    phase = torch.atan2(spec[...,1], spec[...,0]) # calculating the piecewise angle of the complex matrix, this is the same as np.angle\n",
    "    mag = mag[...,1:-1]       # resulting shape is 256x258 --> cut down the edges to have 256x256 square shaped matrices\n",
    "    phase = phase[...,1:-1]\n",
    "\n",
    "    \n",
    "    return mag, phase\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nnIT_WxlGNwH"
   },
   "outputs": [],
   "source": [
    "# heatmap calculation of the magnitude to obtain visible spectrograms\n",
    "def heatmap(mag):\n",
    "    mag = np.log10(mag+1)     # bringing the magnitude to the visible spectrum\n",
    "    for i in range(len(mag)):         #normalizing the values of the magnitude within the batch\n",
    "      mag[i] = mag[i]/np.max(mag[i])\n",
    "    mag = mag.astype(np.float32)    \n",
    "    mag=torch.from_numpy(mag.copy()).unsqueeze(1)\n",
    "\n",
    "    return mag\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cKmJl8C13oxL"
   },
   "outputs": [],
   "source": [
    "#get_batch function used to produce the input spectrograms for the network from the audio clips\n",
    "def get_batch(batch_size,N_=1,instruments_wave=instruments_wave,instrument='random'):\n",
    "\n",
    "    index_list = [[0,1],[1,3],[2,1],[3,1/3]]  # 1st index represents the intrument (bongo,guira,bass,guitar), 2nd index their 'weights' which we use to multiply them to make them sound louder or queieter\n",
    "    instruments = ['bongo','guira','bass','guitar']\n",
    "\n",
    "    #selecting the target instrument. it could be random or specified\n",
    "    if instrument=='random' or instrument == \"all\": \n",
    "      index_pick = random.choice(index_list)\n",
    "    elif instrument=='bongo':\n",
    "      index_pick = index_list[0]\n",
    "    elif instrument=='guira':\n",
    "      index_pick = index_list[1]\n",
    "    elif instrument=='bass':\n",
    "      index_pick = index_list[2]\n",
    "    elif instrument=='guitar':\n",
    "      index_pick = index_list[3]\n",
    "\n",
    "    index = index_pick[0]\n",
    "    index_scale = index_pick[1]\n",
    "\n",
    "    #selecting 2 distinct instances of the target instrument: one target and one true\n",
    "    x_target = torch.stack(random.choices(instruments_wave[index],k=batch_size))\n",
    "    x_true = torch.stack(random.choices(instruments_wave[index],k=batch_size))\n",
    "    for i in range(batch_size):           \n",
    "      while torch.equal(x_true[i],x_target[i]):\n",
    "        x_true[i] = torch.stack(random.choices(instruments_wave[index],k=1))\n",
    "\n",
    "    #remove the already selected instrument so that it won't appear twice in the mixture\n",
    "    index_list.remove(index_pick)\n",
    "   \n",
    "    #choosing the rest of the instruments for the mixture at random\n",
    "    audio_target = (x_target) * index_scale     \n",
    "    audio0 = (x_true) * index_scale     \n",
    "    audio1 = (torch.stack(random.choices(instruments_wave[index_list[0][0]],k=batch_size))) * index_list[0][1]\n",
    "    audio2 = (torch.stack(random.choices(instruments_wave[index_list[1][0]],k=batch_size))) * index_list[1][1]\n",
    "    audio3 = (torch.stack(random.choices(instruments_wave[index_list[2][0]],k=batch_size))) * index_list[2][1]\n",
    "    \n",
    "    audios = [audio0, audio1, audio2, audio3] \n",
    "    \n",
    "    #adding the audios together to create the mixture\n",
    "    wave_mix = audio0 + audio1 + audio2 + audio3\n",
    "\n",
    "    #deriving the magnitude (or amplitude) and phase matrices of the target, the ground truth and the mixture audios\n",
    "    #and producing the spectrograms from the amplitude matrices using the heatmap\n",
    "    amp_mix, phase_mix = _stft(wave_mix)\n",
    "    spec_mix = heatmap(amp_mix.numpy())\n",
    "\n",
    "    amp_true,phase_true = _stft(audio0)\n",
    "    spec_true= heatmap(amp_true.numpy())\n",
    "\n",
    "    amp_target,phase_target = _stft(audio_target)\n",
    "    spec_target = heatmap(amp_target.numpy())\n",
    "\n",
    "    #this section is only used when extracting all 4 instruments from the mixture during testing\n",
    "    if instrument == 'all':\n",
    "      target_0 = torch.stack(random.choices(instruments_wave[index],k=batch_size))\n",
    "      target_1 = torch.stack(random.choices(instruments_wave[index_list[0][0]],k=batch_size))\n",
    "      target_2 = torch.stack(random.choices(instruments_wave[index_list[1][0]],k=batch_size))\n",
    "      target_3 = torch.stack(random.choices(instruments_wave[index_list[2][0]],k=batch_size))\n",
    "      \n",
    "      #creating a target for each instrument, note that no ground truth needed at this point\n",
    "      for i in range(batch_size):\n",
    "        while torch.equal(audio0[i],target_0[i]):\n",
    "          target_0[i] = torch.stack(random.choices(instruments_wave[index],k=1))\n",
    "        while torch.equal(audio1[i],target_1[i]):\n",
    "          target_1[i] = torch.stack(random.choices(instruments_wave[index_list[0][0]],k=1))\n",
    "        while torch.equal(audio2[i],target_2[i]):\n",
    "          target_2[i] = torch.stack(random.choices(instruments_wave[index_list[1][0]],k=1))\n",
    "        while torch.equal(audio3[i],target_3[i]):\n",
    "          target_3[i] = torch.stack(random.choices(instruments_wave[index_list[2][0]],k=1))\n",
    "      \n",
    "      #producing the spectrograms for each target instances\n",
    "      amp_target0,_ = _stft(target_0)\n",
    "      spec_target0 = heatmap(amp_target0.numpy())\n",
    "      amp_target1,_ = _stft(target_1)\n",
    "      spec_target1 = heatmap(amp_target1.numpy())\n",
    "      amp_target2,_ = _stft(target_2)\n",
    "      spec_target2 = heatmap(amp_target2.numpy())\n",
    "      amp_target3,_ = _stft(target_3)\n",
    "      spec_target3 = heatmap(amp_target3.numpy())\n",
    "\n",
    "      #the order at which the targets were picked\n",
    "      instrument_list=[instruments[index],instruments[index_list[0][0]],instruments[index_list[1][0]],instruments[index_list[2][0]]]\n",
    "      #use this only during testing\n",
    "      return spec_mix,phase_mix,wave_mix,audio0,spec_target0,audio1,spec_target1,audio2,spec_target2,audio3,spec_target3,instrument_list\n",
    "\n",
    "    #use this during training\n",
    "    return phase_mix, phase_true, phase_target, wave_mix, spec_mix, spec_true, spec_target, amp_mix, instruments[index], audio0\n",
    "\n",
    "bsize = 4\n",
    "phase_mix,phase_true, phase_target, audio_mix, spec_mix, spec_true, spec_target, ori_amp_mix, instrument,true_original  = get_batch(bsize)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 250
    },
    "colab_type": "code",
    "id": "devnbHRJnFcj",
    "outputId": "7aff18a6-2f0f-47f0-ac54-d5c4694e18fe"
   },
   "outputs": [],
   "source": [
    "#demonstrating a random spectrogram of each instrument\n",
    "specmix,_,_,_,spec_target0,_,spec_target1,_,spec_target2,_,spec_target3,inst_list=get_batch(bsize,instrument='all')\n",
    "fig=plt.figure(figsize=(20,5))\n",
    "fig.add_subplot(1,5,1).set_title('mix',fontsize=18)\n",
    "plt.imshow(specmix[0][0])\n",
    "fig.add_subplot(1,5,2).set_title(inst_list[0],fontsize=18)\n",
    "plt.imshow(spec_target0[0][0])\n",
    "fig.add_subplot(1,5,3).set_title(inst_list[1],fontsize=18)\n",
    "plt.imshow(spec_target1[0][0])\n",
    "fig.add_subplot(1,5,4).set_title(inst_list[2],fontsize=18)\n",
    "plt.imshow(spec_target2[0][0])\n",
    "fig.add_subplot(1,5,5).set_title(inst_list[3],fontsize=18)\n",
    "plt.imshow(spec_target3[0][0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 713
    },
    "colab_type": "code",
    "id": "2DA2PquijIld",
    "outputId": "5c7e4d34-e1f5-4a01-e32c-882aa866e1f8"
   },
   "outputs": [],
   "source": [
    "#demonstrating the target, ground truth and mixture \n",
    "fig = plt.figure(figsize=(3*bsize,5*2))\n",
    "for b in range(bsize):\n",
    "  plt.subplot(3,bsize,b+1).set_title('spec mix {0}'.format(b))\n",
    "  plt.imshow(spec_mix[b,0].numpy())\n",
    "\n",
    "  plt.subplot(3,bsize,b+1+bsize).set_title('spec target {0}'.format(b))\n",
    "  plt.imshow(spec_target[b,0].numpy())\n",
    "\n",
    "  plt.subplot(3,bsize,b+1+2*bsize).set_title('spec true {0}'.format(b))\n",
    "  plt.imshow(spec_true[b,0].numpy())\n",
    "plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "13tTZz_D2pnD",
    "outputId": "7071e8fc-61a5-448b-d57a-6a45f0023d6c"
   },
   "outputs": [],
   "source": [
    "# defining our U-Net architcture\n",
    "class UNet(nn.Module):\n",
    "    def __init__(self, in_class=2, out_class=1, test = -1):     #two in class: mixture,target; one out class: prediction\n",
    "        super().__init__()\n",
    "        self.instrument_index = test\n",
    "                \n",
    "        self.dconv_down1 = self.conv_block(in_class, 64)\n",
    "        self.dconv_down2 = self.conv_block(64, 128)\n",
    "        self.dconv_down3 = self.conv_block(128, 256)\n",
    "        self.dconv_down4 = self.conv_block(256, 512)\n",
    "        self.dconv_down5 = self.conv_block(512, 1024)\n",
    "\n",
    "        self.dconv_up4 = self.conv_block(512 + 1024, 512)\n",
    "        self.dconv_up3 = self.conv_block(256 + 512, 256)\n",
    "        self.dconv_up2 = self.conv_block(128 + 256, 128)\n",
    "        self.dconv_up1 = self.conv_block(64 + 128, 64)\n",
    "        \n",
    "        self.conv_last = nn.Conv2d(64, out_class, 1)\n",
    "\n",
    "\n",
    "    def conv_block(self, in_channels, out_channels):                #convolutional block we used in down and upsampling\n",
    "        return nn.Sequential(\n",
    "            nn.Conv2d(in_channels, out_channels, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.ReLU(inplace=True)\n",
    "    )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        z = self.encode(x)\n",
    "        out = self.decode(z)\n",
    "        return out\n",
    "\n",
    "    def encode(self,x):             #downsampling\n",
    "        self.conv1 = self.dconv_down1(x) # skip-connection 1\n",
    "        x = F.max_pool2d(self.conv1, 2)\n",
    "\n",
    "        self.conv2 = self.dconv_down2(x) # skip-connection 2\n",
    "        x = F.max_pool2d(self.conv2, 2)\n",
    "\n",
    "        self.conv3 = self.dconv_down3(x) # skip-connection 3\n",
    "        x = F.max_pool2d(self.conv3, 2)\n",
    "\n",
    "        self.conv4 = self.dconv_down4(x) # skip-connection 4\n",
    "        x = F.max_pool2d(self.conv4, 2)\n",
    "\n",
    "        x = self.dconv_down5(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def decode(self,x):             #upsampling\n",
    "        x = F.interpolate(x, scale_factor=2, mode='bilinear')    \n",
    "        x = torch.cat([x, self.conv4], dim=1)    # skip-connection 4\n",
    "        x = self.dconv_up4(x)\n",
    "\n",
    "        x = F.interpolate(x, scale_factor=2, mode='bilinear')    \n",
    "        x = torch.cat([x, self.conv3], dim=1)    # skip-connection 3\n",
    "        x = self.dconv_up3(x)\n",
    "\n",
    "        x = F.interpolate(x, scale_factor=2, mode='bilinear')\n",
    "        x = torch.cat([x, self.conv2], dim=1)    # skip-connection 2\n",
    "        x = self.dconv_up2(x)\n",
    "\n",
    "        x = F.interpolate(x, scale_factor=2, mode='bilinear')\n",
    "        x = torch.cat([x, self.conv1], dim=1)    # skip-connection 1\n",
    "        x = self.dconv_up1(x)\n",
    "\n",
    "        x = self.conv_last(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "N = UNet().to(device)\n",
    "\n",
    "print(f'> Number of network parameters {len(torch.nn.utils.parameters_to_vector(N.parameters()))}')\n",
    "\n",
    "# initialise the optimiser\n",
    "optimiser = torch.optim.Adam(N.parameters(), lr=0.001)\n",
    "epoch = 0\n",
    "liveplot = PlotLosses()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "h1aP9k5S12pB"
   },
   "outputs": [],
   "source": [
    "#load previously saved trained models\n",
    "def loadmodel(model,path):\n",
    "  checkpoint = torch.load(path)\n",
    "  model.load_state_dict(checkpoint['model_state_dict'])\n",
    "  optimiser.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "  epoch = checkpoint['epoch']\n",
    "  loss = checkpoint['loss']\n",
    "  loss_list = checkpoint['loss_list']\n",
    "  \n",
    "  model.train()\n",
    "  return model,optimiser,epoch,loss,loss_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 442
    },
    "colab_type": "code",
    "id": "Co265Fjard7C",
    "outputId": "d2b72c89-5008-40e4-8efe-c145308f806d"
   },
   "outputs": [],
   "source": [
    "#defining path for model to be saved during training\n",
    "model_save_name = 'test_model.pt'\n",
    "path = \"drive/My Drive/Colab Notebooks/{0}\".format(model_save_name)\n",
    "!ls drive/My\\ Drive/Colab\\ Notebooks/*pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 729
    },
    "colab_type": "code",
    "id": "0MhtRPJa3Ity",
    "outputId": "c25577ff-cd89-458b-abce-ed6ccb58cd91"
   },
   "outputs": [],
   "source": [
    "#model training\n",
    "loss_list_=[]\n",
    "save_bool = True    #indicate to save model whilst training at every epoch\n",
    "K = 3 #number of prediction samples to show at each epoch\n",
    "while epoch<2:\n",
    "  print('epoch',epoch)\n",
    "  logs = {}\n",
    "  train_loss_arr = np.zeros(0)\n",
    "  for i in range(100):\n",
    "    phase_mix,phase_true, phase_target,_, x_mix,x_true,x_target,_,instrument,_ = get_batch(16) \n",
    "    x_target,x_true,x_mix = x_target.to(device), x_true.to(device), x_mix.to(device) \n",
    "    optimiser.zero_grad()\n",
    "    cnn_in = torch.cat([x_target,x_mix], dim=1).to(device)    #constructing the input for the network of size [Bx2x256x256]\n",
    "\n",
    "    p = N(cnn_in)             #prediction\n",
    "    p = torch.clamp(p,0,1)    #clipping values between 0 and 1\n",
    "\n",
    "    loss = ((p-x_true)**2).mean()     #mean squared error loss between predicted and ground truth spectrograms\n",
    "    loss.backward()\n",
    "    optimiser.step()\n",
    "    # torch.cuda.empty_cache()\n",
    "  \n",
    "  #prediction demonstration at every epoch \n",
    "  clear_output(wait=True)\n",
    "  fig = plt.figure(figsize=(15,6))\n",
    "  fig.add_subplot(2,2,1).set_title('mix epoch {0}'.format(epoch))\n",
    "  plt.imshow(torchvision.utils.make_grid(x_mix[:K]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))  \n",
    "  fig.add_subplot(2,2,3).set_title('target {0}'.format(instrument))\n",
    "  plt.imshow(torchvision.utils.make_grid(x_target[:K]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))\n",
    "  fig.add_subplot(2,2,2).set_title('true {0}'.format(instrument))\n",
    "  plt.imshow(torchvision.utils.make_grid(x_true[:K]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))\n",
    "  fig.add_subplot(2,2,4).set_title('predicted {0}'.format(instrument))\n",
    "  plt.imshow(torchvision.utils.make_grid(p[:K]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))  \n",
    "  \n",
    "\n",
    "  train_loss_arr = np.append(train_loss_arr, loss.cpu().data)\n",
    "  liveplot.update({\n",
    "        'loss': train_loss_arr.mean()\n",
    "  })\n",
    "  loss_list_.append(train_loss_arr)   #collecting the losses to save them\n",
    "\n",
    "  liveplot.draw()     #plotting the loss\n",
    "  plt.tight_layout()\n",
    "  plt.show()\n",
    "  epoch = epoch+1\n",
    "  \n",
    "  #saving the state of the model at every epoch\n",
    "  if save_bool:\n",
    "    torch.save({\n",
    "              'epoch': epoch,\n",
    "              'model_state_dict': N.state_dict(),\n",
    "              'optimizer_state_dict': optimiser.state_dict(),\n",
    "              'loss': loss,\n",
    "              'loss_list': loss_list_,\n",
    "              }, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xW-LE11rAXii"
   },
   "outputs": [],
   "source": [
    "#function to plot the average loss over every x epoch\n",
    "def plotmeanloss(loss_list,epoch_range_,fname=None):\n",
    "  max_epoch = len(loss_list)\n",
    "  epoch_range = [x for x in range(max_epoch) if x % epoch_range_ == 0]\n",
    "  y=[]\n",
    "  x=[]\n",
    "  y_err=[]\n",
    "  for i in range(len(epoch_range)):\n",
    "    try:      #check if total epochs mod x is 0\n",
    "      y_range = loss_list[epoch_range[i]:epoch_range[i+1]]\n",
    "      y.append(np.mean(y_range))  \n",
    "      y_err.append(np.std(y_range))     #error calculation\n",
    "      x.append(epoch_range[i]+epoch_range_)\n",
    "    except:           #if the total number of epochs mod x is not 0, then take the average over the remaining epochs at the end\n",
    "      y_range = loss_list[epoch_range[i]:]\n",
    "      x.append(max_epoch)\n",
    "      y.append(np.mean(y_range))        \n",
    "      y_err.append(np.std(y_range))      #error calculation\n",
    "  fig = plt.figure(figsize=(6,5))\n",
    "  if fname!=None:\n",
    "    fig.suptitle('{0}'.format(fname))\n",
    "  plt.title('Loss over every {0} epochs'.format(epoch_range_))\n",
    "  if epoch_range_ ==1:    #if x==1, then don't plot error bars\n",
    "    plt.plot(x,y)\n",
    "  else:                   # if x>1, do an errorbar plot with the calculated errors\n",
    "    plt.errorbar(x,y,y_err,0,fmt='o-',mew=2,ms=7,ecolor='gray',capsize=5,elinewidth=2)\n",
    "  plt.xlabel('epoch')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pYEgGhkJqY8-"
   },
   "outputs": [],
   "source": [
    "#reversing the previously seen heatmap calculations to obtain the predicted magnitude matrix\n",
    "def inverse_heatmap(heatmap_):\n",
    "  mag1 = heatmap_.squeeze(1).numpy()\n",
    "  mag2 = 10**mag1             #reversing the log_10\n",
    "  mag3 = mag2-1\n",
    "  mag = torch.from_numpy(mag3)\n",
    "  return mag\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Gv_4d-xXtfeH"
   },
   "outputs": [],
   "source": [
    "#reversing the magnitude and phase matrix calculations to produce an input for the inverse STFT\n",
    "def istft_reconstruction(mag, phase):\n",
    "\n",
    "  # real and imaginary parts from the angles\n",
    "  phase_real = torch.cos(phase)     \n",
    "  phase_imag = torch.sin(phase)\n",
    "\n",
    "  #constructing the complex matrix input for the istft from the magnitude and the real and imaginary parts obtained from the phase\n",
    "  complex_matrix = torch.stack((mag*phase_real,mag*phase_imag),dim=-1)\n",
    "\n",
    "  #producing the predicted audio from the complex matrix by applying inverse STFT calculations\n",
    "  wav = torchaudio.functional.istft(complex_matrix,n_fft=510,hop_length=70,win_length=510,window=torch.hann_window(510),length=18000) #fixing the sample length to be 18000 preserved the audio duration\n",
    "  return wav"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "40547055abb44281ba651874f77ea27e",
      "41b39f84abbe4c2597ff6f78a4dcb63b",
      "8eaca2816ccb4e11a6f3093109718aae"
     ]
    },
    "colab_type": "code",
    "id": "yWAg0VCXXNnf",
    "outputId": "61b8ddcc-e8e1-40ab-ab79-63172b681bab"
   },
   "outputs": [],
   "source": [
    "#choose which trained model to load\n",
    "modeldict = 'drive/My Drive/Colab Notebooks/'\n",
    "modellist=[]\n",
    "for filename in os.listdir(modeldict):\n",
    "  if filename.endswith('.pt'):\n",
    "    modellist.append(filename)\n",
    "import ipywidgets as widgets\n",
    "model_picker = widgets.Dropdown(options=modellist)\n",
    "model_picker\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "nkIiuWOBaXwi",
    "outputId": "e11f8da0-722b-4ac8-975b-8cc0a2a5acb7"
   },
   "outputs": [],
   "source": [
    "#define the network and load the selected model into it\n",
    "N = UNet().to(device)\n",
    "N,_,_,_,loss_list = loadmodel(N,'{0}{1}'.format(modeldict,model_picker.value))\n",
    "print(model_picker.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 350
    },
    "colab_type": "code",
    "id": "j28KlWSW9Itk",
    "outputId": "1ee55eff-06c3-4629-f556-249f4d5b6daa"
   },
   "outputs": [],
   "source": [
    "plotmeanloss(loss_list,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AOZTlo4VoS_E"
   },
   "outputs": [],
   "source": [
    "#testing the trained model\n",
    "def test(model,myinstrument,bsize,show_audio=True):\n",
    "\n",
    "  if myinstrument == \"all\":     #extracting all instruments from the test clip\n",
    "    prediction_list=[]\n",
    "    if bsize == 1:    # minimum batch size of 2 is required for the get_batch to work, however later on we just going to consider the 1st batch and ignore the second one\n",
    "      my_x_mix, my_phase_mix,audio_mix, true0, target0, true1, target1, true2, target2, true3, target3, instrument_list = get_batch(2,instrument=myinstrument)\n",
    "    else:   #load mix, target and ground truth\n",
    "      my_x_mix, my_phase_mix,audio_mix, true0, target0, true1, target1, true2, target2, true3, target3, instrument_list = get_batch(bsize,instrument=myinstrument)\n",
    "    target_list=[target0,target1,target2,target3]\n",
    "    true_list = [true0,true1,true2,true3]\n",
    "\n",
    "    # for b in range(bsize):      #display the mixture audio\n",
    "    #   print('     MIX example {0}'.format(b))\n",
    "    #   display(Audio(audio_mix[b],rate=11025))\n",
    "\n",
    "    for target in range(4):   #for each instrument, take the correspoinding spectrograms and feed into the network\n",
    "      print('\\nINSTRUMENT {0}\\n'.format(instrument_list[target]))\n",
    "      cnn_in = torch.cat([target_list[target],my_x_mix], dim=1).to(device)  \n",
    "      prediction = model(cnn_in)\n",
    "      prediction = torch.clamp(prediction,0,1)  #as before, clamp predictions between 0 and 1\n",
    "      amp_inverse_prediction = inverse_heatmap(prediction.detach().cpu())     #predicted magnitude\n",
    "      wav_prediction = istft_reconstruction(amp_inverse_prediction,my_phase_mix)      #predicted audio\n",
    "      wav_prediction = wav_prediction.cpu().numpy()\n",
    "      prediction_list.append(wav_prediction)\n",
    "\n",
    "      for b in range(bsize):      #display the mix and the predicted audio together with the correspongind ground truth audio for easy comparison\n",
    "        print('example {0}'.format(b))\n",
    "        # print(true_list[target][b].shape)\n",
    "        print('MIX'.format(b))\n",
    "        display(Audio(audio_mix[b],rate=11025))\n",
    "        print('prediction')\n",
    "        display(Audio(wav_prediction[b],rate=11025))\n",
    "        print('true')\n",
    "        display(Audio(true_list[target][b],rate=11025))\n",
    "    return prediction_list,true_list,audio_mix,instrument_list\n",
    "\n",
    "  else:     #if a target instrument is specified, only fetch a target spectrogram of that instrument\n",
    "    my_phase_mix,my_phase_true, my_phase_target,_, my_x_mix,my_x_true,my_x_target,_,my_instrument,true_original = get_batch(bsize,instrument=myinstrument)\n",
    "  \n",
    "    cnn_in = torch.cat([my_x_target,my_x_mix], dim=1).to(device)    # as before, input to our cnn\n",
    "    prediction = model(cnn_in)\n",
    "    prediction = torch.clamp(prediction,0,1)    #predictions clamped as before\n",
    "\n",
    "    #plot the mixture, the target, the ground truth, and the predicted spectrograms together for easy comparison\n",
    "    fig = plt.figure(figsize=(15,8))\n",
    "    fig.add_subplot(3,2,1).set_title('mix')\n",
    "    plt.imshow(torchvision.utils.make_grid(my_x_mix[:]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))  \n",
    "    fig.add_subplot(3,2,3).set_title('target {0}'.format(my_instrument))\n",
    "    plt.imshow(torchvision.utils.make_grid(my_x_target[:]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))\n",
    "    fig.add_subplot(3,2,2).set_title('true {0}'.format(my_instrument))\n",
    "    plt.imshow(torchvision.utils.make_grid(my_x_true[:]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))\n",
    "    fig.add_subplot(3,2,4).set_title('predicted {0}'.format(my_instrument))\n",
    "    plt.imshow(torchvision.utils.make_grid(prediction[:]).cpu().data.permute(0,2,1).contiguous().permute(2,1,0))\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    #producing the magnitudes and audio of the predicted, ground truth, and mixture spectrograms\n",
    "    amp_inverse_prediction = inverse_heatmap(prediction.detach().cpu())\n",
    "    wav_prediction = istft_reconstruction(amp_inverse_prediction,my_phase_mix)\n",
    "    wav_prediction = wav_prediction.cpu().numpy()\n",
    "\n",
    "    amp_inverse_true = inverse_heatmap(my_x_true.detach().cpu())\n",
    "    wav_true = istft_reconstruction(amp_inverse_true,my_phase_true)\n",
    "    wav_true = wav_true.cpu().numpy()\n",
    "\n",
    "    amp_inverse_mix = inverse_heatmap(my_x_mix.detach().cpu())\n",
    "    wav_mix = istft_reconstruction(amp_inverse_mix,my_phase_mix)\n",
    "    wav_mix = wav_mix.cpu().numpy()\n",
    "\n",
    "    print(my_instrument,'\\n')\n",
    "    #once again display the audio clips together within a batch for easy comparison\n",
    "    if show_audio:\n",
    "      for b in range(bsize):\n",
    "        print('batch {0}\\n'.format(b+1))\n",
    "        print('     PREDICTED')\n",
    "        display(Audio(wav_prediction[b],rate=11025))\n",
    "        print('     TRUE')\n",
    "        display(Audio(wav_true[b],rate=11025))\n",
    "        print('     MIX')\n",
    "        display(Audio(wav_mix[b],rate=11025))\n",
    "        print('\\n\\n')\n",
    "    return wav_prediction,wav_true,wav_mix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WjkYuNfdo3ty"
   },
   "outputs": [],
   "source": [
    "#choosing which instrument to filter and how many samples(batch size) we want\n",
    "\n",
    "instrument_to_extract = \"all\" #@param [\"random\",\"bongo\", \"guira\", \"bass\", \"guitar\",'all']\n",
    "\n",
    "batch_size =  5#@param {type:\"integer\"}\n",
    "\n",
    "preds,trues,mix, inst_list = test(N,instrument_to_extract,batch_size,show_audio=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cSmxvqsZCVNK"
   },
   "outputs": [],
   "source": [
    "# this stretching function was obtained from https://github.com/gaganbahga/time_stretch\n",
    "def stretch(x, factor,length=0, nfft=2048):\n",
    "    '''\n",
    "    stretch an audio sequence by a factor using FFT of size nfft converting to frequency domain\n",
    "    :param x: np.ndarray, audio array in PCM float32 format\n",
    "    :param factor: float, stretching or shrinking factor, depending on if its > or < 1 respectively\n",
    "    :return: np.ndarray, time stretched audio\n",
    "    '''\n",
    "    stft = librosa.core.stft(x, n_fft=nfft).transpose()  # i prefer time-major fashion, so transpose\n",
    "    stft_rows = stft.shape[0]\n",
    "    stft_cols = stft.shape[1]\n",
    "\n",
    "    times = np.arange(0, stft.shape[0], factor)  # times at which new FFT to be calculated\n",
    "    hop = nfft/4                                 # frame shift\n",
    "    stft_new = np.zeros((len(times), stft_cols), dtype=np.complex_)\n",
    "    phase_adv = (2 * np.pi * hop * np.arange(0, stft_cols))/ nfft\n",
    "    phase = np.angle(stft[0])\n",
    "\n",
    "    stft = np.concatenate( (stft, np.zeros((1, stft_cols))), axis=0)\n",
    "\n",
    "    for i, time in enumerate(times):\n",
    "        left_frame = int(np.floor(time))\n",
    "        local_frames = stft[[left_frame, left_frame + 1], :]\n",
    "        right_wt = time - np.floor(time)                        # weight on right frame out of 2\n",
    "        local_mag = (1 - right_wt) * np.absolute(local_frames[0, :]) + right_wt * np.absolute(local_frames[1, :])\n",
    "        local_dphi = np.angle(local_frames[1, :]) - np.angle(local_frames[0, :]) - phase_adv\n",
    "        local_dphi = local_dphi - 2 * np.pi * np.floor(local_dphi/(2 * np.pi))\n",
    "        stft_new[i, :] =  local_mag * np.exp(phase*1j)\n",
    "        phase += local_dphi + phase_adv\n",
    "\n",
    " \n",
    "    out = librosa.core.istft(stft_new.transpose())\n",
    "    return out, x.shape[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KUAR8glVT0cB"
   },
   "outputs": [],
   "source": [
    "def extraction(model,mysong,instrument):\n",
    "  '''\n",
    "  extract \"instrument\" from \"mysong\" using the \"model\" neural network\n",
    "  '''\n",
    "\n",
    "  tempo_ori,beats = librosa.beat.beat_track(mysong.numpy(),sr=44100,start_bpm=100.0,units=\"samples\",trim=True,tightness=100)  #identify beat locations\n",
    "  out = []\n",
    "  i=0\n",
    "  bar = 4     # we are interested in 4-beat-long clips\n",
    "  while i+bar < len(beats):\n",
    "    # while test_instrument != instrument:\n",
    "    _,_, _,_, _,_,test_x_target,_,test_instrument,_ = get_batch(2,instrument=instrument)    #fetch a target instrument clip from training data\n",
    "    test_x_target = test_x_target[0].unsqueeze(0)\n",
    "    \n",
    "    song_sample = mysong[beats[i]:beats[i+bar]]   #generate clips from the song every 4 beats\n",
    "    sample_resampled = torchaudio.transforms.Resample(44100, new_freq=11025)(song_sample.view(1,-1))[0]   #resample clips to have the same sample rate as the training data\n",
    "\n",
    "    sample_length = sample_resampled.shape[0]   #sample length of the resampled clip\n",
    "   \n",
    "    shrinkrate = sample_length/18000      #stretching factor\n",
    "    \n",
    "    nfft = 1000 #window size of the fast Fourier transform used in stretching\n",
    "    new_sample, inverse_shape = stretch(sample_resampled.numpy(),shrinkrate,nfft=nfft)  #stretch the clip so that it has the same tempo and hence sample length as the training clips. also save the length of the original clip\n",
    "    new_sample_tensor = torch.from_numpy(new_sample).unsqueeze(0)\n",
    "    sample_amp, sample_phase = _stft(new_sample_tensor)       #produce the magnitude and phase matrices of size 256x256 of the clip \n",
    "    sample_heatmap = heatmap(sample_amp.unsqueeze(0).numpy())   #compute spectrogram\n",
    "    sample_cnn_in = torch.cat([test_x_target,sample_heatmap], dim=1).to(device) #input of the model using the spectrogram of the clip from the song and the target one from the training dataset\n",
    "    sample_p = model(sample_cnn_in)     #predicted spectrogram\n",
    "    sample_p = torch.clamp(sample_p,0,1)    #clamping prediction\n",
    "\n",
    "    amp_inverse = inverse_heatmap(sample_p.detach().cpu())      #retrieve predicted magnitude\n",
    "    wav = istft_reconstruction(amp_inverse,sample_phase)        #construct predicted audio\n",
    "    wav = wav.cpu().numpy()[0]\n",
    "    wav,_ = stretch(wav,1/shrinkrate,nfft = 512, length=inverse_shape)    #reverse the stretching of the audio / \"unstretch\" it using the previously saved original sample length of the clip, so that it has the same tempo as the original song\n",
    "    out.extend(wav)       #collect and concatenate the predictions \n",
    "    i+=bar\n",
    "\n",
    "  return np.array(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "eb3a12ab4b7649ff8dd828ed48cfffcb",
      "5c62fde7d56e450f996c811ec79bc374",
      "152b78604f6942599477328392de4e9c"
     ]
    },
    "colab_type": "code",
    "id": "-lLYen8QvVKD",
    "outputId": "006792f2-fa34-4b04-d779-4c7ecd3cd3e8"
   },
   "outputs": [],
   "source": [
    "#pick a song to process\n",
    "songs_dir = \"drive/My Drive/Colab Notebooks/Bachata/songs/\"\n",
    "songlist=[]\n",
    "songnames=[]\n",
    "songdir=sorted([x for x in os.listdir(songs_dir) if x.endswith('.mp3')])\n",
    "for name in songdir:\n",
    "    mysong_ori, sample_rate = torchaudio.load(\"{0}{1}\".format(songs_dir,name))\n",
    "    mysong = mysong_ori[0]\n",
    "    songlist.append(mysong)\n",
    "    songnames.append(name)\n",
    "\n",
    "song_picker = widgets.Dropdown(options=songnames)\n",
    "song_picker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZtP1YkOWy3Sr"
   },
   "outputs": [],
   "source": [
    "def originalsong(name,songlist,songnames):\n",
    "  '''\n",
    "  displaying original song\n",
    "  '''\n",
    "  print(\"original song\",name)\n",
    "  song = songlist[songnames.index(name)]\n",
    "  display(Audio(song,rate=44100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 92
    },
    "colab_type": "code",
    "id": "qO_Cr3UvPzc4",
    "outputId": "5c8f5f45-c0ea-4c98-8b1b-5ea217b120ac"
   },
   "outputs": [],
   "source": [
    "#display original song\n",
    "originalsong(song_picker.value,songlist,songnames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 609
    },
    "colab_type": "code",
    "id": "1gfAQfh_H_Rw",
    "outputId": "8749b24a-42c2-4a89-f7a6-f2c908f64a78"
   },
   "outputs": [],
   "source": [
    "#extract all 4 instruments\n",
    "print(model_picker.value)\n",
    "print(song_picker.value)\n",
    "instrument_list = [\"bongo\", \"guira\", \"bass\", \"guitar\"]\n",
    "song = songlist[songnames.index(song_picker.value)]\n",
    "for i in instrument_list:\n",
    "  wav = extraction(N,song,i)\n",
    "  print(\"\\nextracted\",i)\n",
    "  display(Audio(wav,rate=11025))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1GkhWS2DyshO"
   },
   "outputs": [],
   "source": [
    "#extract a specific instrument\n",
    "instrument = \"guitar\" #@param [\"bongo\", \"guira\", \"bass\", \"guitar\"]\n",
    "song = songlist[songnames.index(song_picker.value)]\n",
    "wav = extraction(N,song,instrument)\n",
    "print(\"extracted\",instrument)\n",
    "display(Audio(wav,rate=11025))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "audiomodel.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.5"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "152b78604f6942599477328392de4e9c": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "40547055abb44281ba651874f77ea27e": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DropdownModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DropdownModel",
      "_options_labels": [
       "model_100_100.pt",
       "model_100_100_noclampspec.pt",
       "mymodel_100_100.pt",
       "mymodel_100_100_normspec.pt",
       "test2_20_50_noscale_N_1.pt",
       "test3_20_50_scale300_N_1.pt",
       "test5_20_50_scale1_N_1_normperbatch.pt",
       "test6_20_50_scale1_N_1_normoverbatch.pt",
       "test7_20_50_scale300_N_1_clip255_normperbatch.pt",
       "test8_20_50_scale200_N_1_clip255_normoverbatch.pt",
       "test9_20_50_scale300_N_1_clip255_normoverbatch.pt",
       "test10_20_50_scale200_N_1_clip255_normperbatch.pt",
       "test11_100_100_scale1_noclip_normperbatch.pt",
       "test12_100_100_scale300_N_1_clip255_normperbatch.pt",
       "DEMO_test11.pt",
       "test13_no_norm.pt",
       "test14_100_100_no_norm.pt",
       "test15_100_100_no_norm.pt",
       "final_model.pt",
       "test17_20_50_mixed_dataset_bass.pt",
       "test18_20_50_mixed_dataset_bongo.pt",
       "test19_20_50_mixed_dataset_guira.pt",
       "test20_20_50_mixed_dataset_guitar.pt",
       "test21_20_50_mixed_dataset_random.pt",
       "test_model.pt"
      ],
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "DropdownView",
      "description": "",
      "description_tooltip": null,
      "disabled": false,
      "index": 18,
      "layout": "IPY_MODEL_8eaca2816ccb4e11a6f3093109718aae",
      "style": "IPY_MODEL_41b39f84abbe4c2597ff6f78a4dcb63b"
     }
    },
    "41b39f84abbe4c2597ff6f78a4dcb63b": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "5c62fde7d56e450f996c811ec79bc374": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "8eaca2816ccb4e11a6f3093109718aae": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "eb3a12ab4b7649ff8dd828ed48cfffcb": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DropdownModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DropdownModel",
      "_options_labels": [
       "FAILED - 24 Mujer Perdoname.mp3",
       "FAILED - A Dónde Va el Amor.mp3",
       "FAILED - Amiga Veneno.mp3",
       "FAILED - Creiste.mp3",
       "FAILED - Eso Es Amor.mp3",
       "FAILED - Loco de Amor.mp3",
       "FAILED - Por Favor.mp3",
       "FAILED - SpMusic presenta Jashel - La Copa Rota - Bachata 2014-2013.mp3",
       "FAILED - Suavemte.mp3",
       "GOOD - El Guardaespaldas....mp3",
       "GOOD - Vanidosa.mp3",
       "SOSO - Hoja en Blanco.mp3",
       "SOSO - Prince Royce - Te Robarte.mp3",
       "SOSO - Vagabundo, borracho & Loco.mp3"
      ],
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "DropdownView",
      "description": "",
      "description_tooltip": null,
      "disabled": false,
      "index": 10,
      "layout": "IPY_MODEL_152b78604f6942599477328392de4e9c",
      "style": "IPY_MODEL_5c62fde7d56e450f996c811ec79bc374"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}