audio_style_transfer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import librosa\n",
    "import os\n",
    "# Change `gpu0` to `cpu` to run on CPU\n",
    "os.environ[\"THEANO_FLAGS\"] = \"mode=FAST_RUN,device=gpu0,floatX=float32\"\n",
    "import numpy as np\n",
    "import scipy\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "import lasagne\n",
    "from lasagne.utils import floatX\n",
    "from IPython.display import Audio, display\n",
    "from lasagne.layers import InputLayer, Conv1DLayer as ConvLayer\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!ls inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load style and content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "CONTENT_FILENAME = \"inputs/imperial.mp3\"\n",
    "STYLE_FILENAME = \"inputs/usa.mp3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "display(Audio(CONTENT_FILENAME))\n",
    "display(Audio(STYLE_FILENAME))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Reads wav file and produces spectrum\n",
    "# Fourier phases are ignored\n",
    "N_FFT = 2048\n",
    "def read_audio_spectum(filename):\n",
    "    x, fs = librosa.load(filename)\n",
    "    S = librosa.stft(x, N_FFT)\n",
    "    p = np.angle(S)\n",
    "    return np.log1p(np.abs(S[np.newaxis,:,:430])), fs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "a_content, fs = read_audio_spectum(CONTENT_FILENAME)\n",
    "a_style, fs = read_audio_spectum(STYLE_FILENAME)\n",
    "\n",
    "N_SAMPLES = a_content.shape[2]\n",
    "N_CHANNELS = a_content.shape[1]\n",
    "a_style = a_style[:, :N_CHANNELS, :N_SAMPLES]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize spectrograms for content and style tracks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 5))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title('Content')\n",
    "plt.imshow(a_content[0,:400,:])\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title('Style')\n",
    "plt.imshow(a_style[0,:400,:])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# During our tests, we discovered that it is essential to use extremely large number of conv filters \n",
    "# In this example we use single convolution with 4096 filters\n",
    "\n",
    "N_FILTERS = 4096\n",
    "inputs = InputLayer((1, N_CHANNELS, N_SAMPLES))\n",
    "conv = ConvLayer(inputs, N_FILTERS, 11, W=lasagne.init.GlorotNormal(gain='relu'))\n",
    "\n",
    "# Implementation of losses and optimization is based on artistic style transfer example in lasagne recipes\n",
    "# https://github.com/Lasagne/Recipes/blob/master/examples/styletransfer/Art%20Style%20Transfer.ipynb\n",
    "def gram_matrix(x):\n",
    "    g = T.tensordot(x, x, axes=([2], [2])) / x.shape[2]\n",
    "    return g\n",
    "\n",
    "def style_loss(A, X,):\n",
    "    G1 = gram_matrix(A)\n",
    "    G2 = gram_matrix(X) \n",
    "    loss = ((G1 - G2)**2).sum()\n",
    "    return loss\n",
    "\n",
    "def content_loss(A, X):\n",
    "    return ((A - X)**2).sum()\n",
    "\n",
    "t = np.zeros_like(a_content)\n",
    "\n",
    "content_features = lasagne.layers.get_output(conv, a_content)\n",
    "style_features = lasagne.layers.get_output(conv, a_style)\n",
    "\n",
    "generated = T.tensor3()\n",
    "gen_features = lasagne.layers.get_output(conv, generated)\n",
    "\n",
    "# set ALPHA=1e-3 for more style, or ALPHA=0 to turn off content entirely\n",
    "ALPHA = 1e-2\n",
    "loss = style_loss(style_features, gen_features) +\\\n",
    "            ALPHA * content_loss(content_features, gen_features)\n",
    "grad = T.grad(loss, generated)\n",
    "\n",
    "f_loss = theano.function([generated], loss)\n",
    "f_grad = theano.function([generated], grad)\n",
    "\n",
    "def eval_loss(x0):\n",
    "    x0 = floatX(x0.reshape((1, N_CHANNELS, N_SAMPLES)))\n",
    "    return f_loss(x0).astype('float64')\n",
    "\n",
    "def eval_grad(x0):\n",
    "    x0 = floatX(x0.reshape((1, N_CHANNELS, N_SAMPLES)))\n",
    "    return np.array(f_grad(x0)).flatten().astype('float64')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#initialization with zeros or gaussian noise can be used\n",
    "#zeros don't work with ALPHA=0\n",
    "#t = floatX(np.random.randn(1, N_CHANNELS, N_SAMPLES))\n",
    "t = floatX(np.zeros((1, N_CHANNELS, N_SAMPLES)))\n",
    "\n",
    "res = scipy.optimize.fmin_l_bfgs_b(eval_loss, t.flatten(), fprime=eval_grad, maxfun=500)\n",
    "t = res[0].reshape((1, N_CHANNELS, N_SAMPLES))\n",
    "print res[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invert spectrogram and save the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a = np.zeros_like(a_content[0])\n",
    "a[:N_CHANNELS,:] = np.exp(t[0]) - 1\n",
    "\n",
    "# This code is supposed to do phase reconstruction\n",
    "p = 2 * np.pi * np.random.random_sample(a.shape) - np.pi\n",
    "for i in range(500):\n",
    "    S = a * np.exp(1j*p)\n",
    "    x = librosa.istft(S)\n",
    "    p = np.angle(librosa.stft(x, N_FFT))\n",
    "\n",
    "OUTPUT_FILENAME = 'outputs/out.wav'\n",
    "librosa.output.write_wav(OUTPUT_FILENAME, x, fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print OUTPUT_FILENAME\n",
    "display(Audio(OUTPUT_FILENAME))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize spectrograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(1,3,1)\n",
    "plt.title('Content')\n",
    "plt.imshow(a_content[0,:400,:])\n",
    "plt.subplot(1,3,2)\n",
    "plt.title('Style')\n",
    "plt.imshow(a_style[0,:400,:])\n",
    "plt.subplot(1,3,3)\n",
    "plt.title('Result')\n",
    "plt.imshow(a[:400,:])\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}