Test GACELA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to test GACELA on your data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First, set up the parameters of the network you trained or downloaded from our pre-trained networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from data.audioLoader import AudioLoader\n",
    "from data.trainDataset import TrainDataset\n",
    "from ganSystem import GANSystem\n",
    "import logging\n",
    "\n",
    "# logging.getLogger().setLevel(logging.DEBUG)  # set root logger to debug\n",
    "\n",
    "\"\"\"Just so logging works...\"\"\"\n",
    "formatter = logging.Formatter('%(name)s:%(levelname)s:%(message)s')\n",
    "console_handler = logging.StreamHandler()\n",
    "console_handler.setLevel(logging.DEBUG)\n",
    "console_handler.setFormatter(formatter)\n",
    "logging.getLogger().addHandler(console_handler)\n",
    "\"\"\"Just so logging works...\"\"\"\n",
    "\n",
    "__author__ = 'Andres'\n",
    "\n",
    "signal_split = [480, 64, 480]\n",
    "md = 32\n",
    "\n",
    "params_stft_discriminator = dict()\n",
    "params_stft_discriminator['stride'] = [2, 2, 2, 2, 2]\n",
    "params_stft_discriminator['nfilter'] = [md, 2 * md, 4 * md, 8 * md, 16 * md]\n",
    "params_stft_discriminator['shape'] = [[5, 5], [5, 5], [5, 5], [5, 5], [5, 5]]\n",
    "params_stft_discriminator['data_size'] = 2\n",
    "\n",
    "params_mel_discriminator = dict()\n",
    "params_mel_discriminator['stride'] = [2, 2, 2, 2, 2]\n",
    "params_mel_discriminator['nfilter'] = [md//4, 2 * md//4, 4 * md//4, 8 * md//4, 16 * md//4]\n",
    "params_mel_discriminator['shape'] = [[5, 5], [5, 5], [5, 5], [5, 5], [5, 5]]\n",
    "params_mel_discriminator['data_size'] = 2\n",
    "\n",
    "params_generator = dict()\n",
    "params_generator['stride'] = [2, 2, 2, 2, 2]\n",
    "params_generator['nfilter'] = [8 * md, 4 * md, 2 * md, md, 1]\n",
    "params_generator['shape'] = [[4, 4], [4, 4], [8, 8], [8, 8], [8, 8]]\n",
    "params_generator['padding'] = [[1, 1], [1, 1], [3, 3], [3, 3], [3, 3]]\n",
    "params_generator['residual_blocks'] = 2\n",
    "\n",
    "params_generator['full'] = 256 * md\n",
    "params_generator['summary'] = True\n",
    "params_generator['data_size'] = 2\n",
    "params_generator['in_conv_shape'] = [16, 2]\n",
    "params_generator['borders'] = dict()\n",
    "params_generator['borders']['nfilter'] = [md, 2 * md, md, md / 2]\n",
    "params_generator['borders']['shape'] = [[5, 5], [5, 5], [5, 5], [5, 5]]\n",
    "params_generator['borders']['stride'] = [2, 2, 2, 2]\n",
    "params_generator['borders']['data_size'] = 2\n",
    "params_generator['borders']['border_scale'] = 1\n",
    "# This does not work because of flipping, border 2 need to be flipped tf.reverse(l, axis=[1]), ask Nathanael\n",
    "params_generator['borders']['width_full'] = None\n",
    "\n",
    "# Optimization parameters inspired from 'Self-Attention Generative Adversarial Networks'\n",
    "# - Spectral normalization GEN DISC\n",
    "# - Batch norm GEN\n",
    "# - TTUR ('GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium')\n",
    "# - ADAM  beta1=0 beta2=0.9, disc lr 0.0004, gen lr 0.0001\n",
    "# - Hinge loss\n",
    "# Parameters are similar to the ones in those papers...\n",
    "# - 'PROGRESSIVE GROWING OF GANS FOR IMPROVED QUALITY, STABILITY, AND VARIATION'\n",
    "# - 'LARGE SCALE GAN TRAINING FOR HIGH FIDELITY NATURAL IMAGE SYNTHESIS'\n",
    "# - 'CGANS WITH PROJECTION DISCRIMINATOR'\n",
    "\n",
    "params_optimization = dict()\n",
    "params_optimization['batch_size'] = 64\n",
    "params_stft_discriminator['batch_size'] = 64\n",
    "params_mel_discriminator['batch_size'] = 64\n",
    "\n",
    "params_optimization['n_critic'] = 1\n",
    "params_optimization['generator'] = dict()\n",
    "params_optimization['generator']['optimizer'] = 'adam'\n",
    "params_optimization['generator']['kwargs'] = [0.5, 0.9]\n",
    "params_optimization['generator']['learning_rate'] = 1e-4\n",
    "params_optimization['discriminator'] = dict()\n",
    "params_optimization['discriminator']['optimizer'] = 'adam'\n",
    "params_optimization['discriminator']['kwargs'] = [0.5, 0.9]\n",
    "params_optimization['discriminator']['learning_rate'] = 1e-4\n",
    "\n",
    "# all parameters\n",
    "params = dict()\n",
    "params['net'] = dict()  # All the parameters for the model\n",
    "params['net']['generator'] = params_generator\n",
    "params['net']['stft_discriminator'] = params_stft_discriminator\n",
    "params['net']['mel_discriminator'] = params_mel_discriminator\n",
    "params['net']['prior_distribution'] = 'gaussian'\n",
    "params['net']['shape'] = [1, 512, 1024]  # Shape of the image\n",
    "params['net']['inpainting'] = dict()\n",
    "params['net']['inpainting']['split'] = signal_split\n",
    "params['net']['gamma_gp'] = 10  # Gradient penalty\n",
    "# params['net']['fs'] = 16000//downscale\n",
    "params['net']['loss_type'] = 'wasserstein'\n",
    "\n",
    "params['optimization'] = params_optimization\n",
    "params['summary_every'] = 250  # Tensorboard summaries every ** iterations\n",
    "params['print_every'] = 50  # Console summaries every ** iterations\n",
    "params['save_every'] = 1000  # Save the model every ** iterations\n",
    "# params['summary_dir'] = os.path.join(global_path, name +'_summary/')\n",
    "# params['save_dir'] = os.path.join(global_path, name + '_checkpoints/')\n",
    "\n",
    "args = dict()\n",
    "args['generator'] = params_generator\n",
    "args['stft_discriminator_count'] = 2\n",
    "args['mel_discriminator_count'] = 3\n",
    "args['stft_discriminator'] = params_stft_discriminator\n",
    "args['mel_discriminator'] = params_mel_discriminator\n",
    "args['borderEncoder'] = params_generator['borders']\n",
    "args['stft_discriminator_in_shape'] = [1, 512, 64]\n",
    "args['mel_discriminator_in_shape'] = [1, 80, 64]\n",
    "args['mel_discriminator_start_powscale'] = 2\n",
    "args['generator_input'] = 1440\n",
    "args['optimizer'] = params_optimization\n",
    "args['split'] = signal_split\n",
    "args['log_interval'] = 100\n",
    "args['spectrogram_shape'] = params['net']['shape']\n",
    "args['gamma_gp'] = params['net']['gamma_gp']\n",
    "args['tensorboard_interval'] = 500\n",
    "args['save_path'] = 'saved_results/'\n",
    "args['experiment_name'] = 'fma_rock'\n",
    "args['save_interval'] = 10000\n",
    "\n",
    "args['fft_length'] = 1024\n",
    "args['fft_hop_size'] = 256\n",
    "args['sampling_rate'] = 22050\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "examples_per_file = 32\n",
    "audioLoader = AudioLoader(args['sampling_rate'], args['fft_length'], args['fft_hop_size'], 50)\n",
    "\n",
    "ganSystem = GANSystem(args)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_at_step = 410000\n",
    "start_at_epoch = 0\n",
    "\n",
    "ganSystem.loadModel(start_at_step, start_at_epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data you want to inpaint and process the spectrograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from data.audioLoader import AudioLoader\n",
    "import numpy as np\n",
    "\n",
    "audioLoader = AudioLoader(22050, 1024, 256, 50)\n",
    "\n",
    "\n",
    "dataFolder = \"rock/\"\n",
    "\n",
    "filenames = Path(dataFolder).rglob('*.mp3')\n",
    "files = {}\n",
    "\n",
    "for filename in filenames:\n",
    "    audio = audioLoader.loadSound(filename)\n",
    "    files[filename] = (audio, audioLoader.computeSpectrogram(audio))\n",
    "    \n",
    "\n",
    "audios = np.zeros([64, 256*1024])\n",
    "spectrograms = np.zeros([64, 1, 512, 1024])\n",
    "\n",
    "filenames = list(files.keys())\n",
    "for index in range(64):\n",
    "    filename = np.random.choice(filenames)\n",
    "    \n",
    "    audio = files[filename][0]\n",
    "    \n",
    "    start = int((len(audio)-1024*256) * np.random.rand()) \n",
    "    audios[index] = files[filename][0][start:start+1024*256]    \n",
    "    spectrograms[index, 0] = files[filename][1][:-1, int(start/256):int(start/256)+1024]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate the gap with GACELA and concatenate it with the borders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "left_borders = torch.from_numpy(spectrograms[:, :, :, :args['split'][0]]).float().to(device)\n",
    "right_borders = torch.from_numpy(spectrograms[:, :, :, args['split'][0] + args['split'][1]:]).float().to(device)\n",
    "print('generate')\n",
    "generated_spectrograms = ganSystem.generateGap([left_borders, right_borders])\n",
    "\n",
    "fake_spectrograms = torch.cat((left_borders, generated_spectrograms, right_borders), 3)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot some of the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(30, 14))\n",
    "for i in range(8):\n",
    "    plt.subplot(2, 4, i+1)\n",
    "\n",
    "    plt.imshow(fake_spectrograms[3*i, 0, :,  200:-200].detach().cpu())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Invert the generated spectrograms using PGHI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tifresi.stft import GaussTruncTF, GaussTF\n",
    "from tifresi.transforms import inv_log_spectrogram, log_spectrogram\n",
    "import numpy as np\n",
    "\n",
    "stft = GaussTruncTF(256, 1024)\n",
    "fake_audios = np.zeros([len(audios), fake_spectrograms.shape[-1]*args['fft_hop_size']])\n",
    "\n",
    "for index, (real_audio, spectrogram) in enumerate(zip(audios, fake_spectrograms)):\n",
    "    unprocessed_spectrogram = inv_log_spectrogram((spectrogram-1)*25).squeeze().detach().cpu().numpy()\n",
    "   \n",
    "    unprocessed_spectrogram = np.concatenate([unprocessed_spectrogram,\n",
    "                                  np.ones_like(unprocessed_spectrogram)[0:1, :]*unprocessed_spectrogram.min()], axis=0) # Fill last column of freqs with zeros    \n",
    "    audio = stft.invert_spectrogram(unprocessed_spectrogram)\n",
    "    fake_audios[index] = audio\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listen to the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Audio\n",
    "\n",
    "for generated_audio_signal in fake_audios:\n",
    "    display(Audio(generated_audio_signal[int(len(generated_audio_signal)*0.3):], rate=22050))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare and save the generated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import numpy as np\n",
    "import librosa \n",
    "\n",
    "from tifresi.phase.modGabPhaseGrad import modgabphasegrad\n",
    "from tifresi.phase.pghi_masked import pghi\n",
    "from tifresi.stft import GaussTruncTF, GaussTF\n",
    "hop_size = 256\n",
    "stft_channels = 1024\n",
    "stft = GaussTruncTF(256, 1024)\n",
    "\n",
    "base_folder_name = \"my_test/\"\n",
    "os.mkdir(base_folder_name)\n",
    "\n",
    "real_folder = base_folder_name + 'real/'\n",
    "fake_folder = base_folder_name + \"GACELA/\"\n",
    "\n",
    "os.mkdir(real_folder)\n",
    "os.mkdir(fake_folder)\n",
    "\n",
    "\n",
    "for index, (real_audio, fake_audio) in enumerate(zip(audios, fake_audios)):\n",
    "    #Real\n",
    "    librosa.output.write_wav(real_folder + str(index) + '.wav', real_audio, sr=22050)\n",
    "\n",
    "    #Fake\n",
    "    librosa.output.write_wav(fake_folder + str(index) + '.wav', fake_audio, sr=22050)\n"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}