Pix2Pix.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pix2Pix in Tensorflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorials walks through an implementation of Pix2Pix as described in [Image-to-Image Translation with Conditional Adversarial Networks](https://arxiv.org/abs/1611.07004).\n",
    "\n",
    "This specific implementation is designed to 'remaster' black-and-white frames from films with 4:3 aspect ratio into full-color and 16:9 aspect ratio frames. \n",
    "\n",
    "To learn more about the Pix2Pix framework, and the images that can be generated using this framework, see my [Medium post](https://medium.com/p/f4d551fa0503).\n",
    "\n",
    "This notebook requires the additional helper.py file, which can be obtained [here]()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#Import the libraries we will need.\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow.contrib.slim as slim\n",
    "import os\n",
    "import scipy.misc\n",
    "import scipy\n",
    "from PIL import Image\n",
    "from glob import glob\n",
    "import os\n",
    "from helper import *\n",
    "%matplotlib inline\n",
    "\n",
    "#Size of image frames\n",
    "height = 144\n",
    "width = 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the Adversarial Networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def generator(c):\n",
    "    with tf.variable_scope('generator'):\n",
    "        #Encoder\n",
    "        enc0 = slim.conv2d(c,64,[3,3],padding=\"SAME\",\n",
    "            biases_initializer=None,activation_fn=lrelu,\n",
    "            weights_initializer=initializer)\n",
    "        enc0 = tf.space_to_depth(enc0,2)\n",
    "        \n",
    "        enc1 = slim.conv2d(enc0,128,[3,3],padding=\"SAME\",\n",
    "            activation_fn=lrelu,normalizer_fn=slim.batch_norm,\n",
    "            weights_initializer=initializer)\n",
    "        enc1 = tf.space_to_depth(enc1,2)\n",
    "\n",
    "        enc2 = slim.conv2d(enc1,128,[3,3],padding=\"SAME\",\n",
    "            normalizer_fn=slim.batch_norm,activation_fn=lrelu,\n",
    "            weights_initializer=initializer)\n",
    "        enc2 = tf.space_to_depth(enc2,2)\n",
    "\n",
    "        enc3 = slim.conv2d(enc2,256,[3,3],padding=\"SAME\",\n",
    "            normalizer_fn=slim.batch_norm,activation_fn=lrelu,\n",
    "            weights_initializer=initializer)\n",
    "        enc3 = tf.space_to_depth(enc3,2)\n",
    "        \n",
    "        #Decoder\n",
    "        gen0 = slim.conv2d(\n",
    "            enc3,num_outputs=256,kernel_size=[3,3],\n",
    "            padding=\"SAME\",normalizer_fn=slim.batch_norm,\n",
    "            activation_fn=tf.nn.elu, weights_initializer=initializer)\n",
    "        gen0 = tf.depth_to_space(gen0,2)\n",
    "\n",
    "        gen1 = slim.conv2d(\n",
    "            tf.concat([gen0,enc2],3),num_outputs=256,kernel_size=[3,3],\n",
    "            padding=\"SAME\",normalizer_fn=slim.batch_norm,\n",
    "            activation_fn=tf.nn.elu,weights_initializer=initializer)\n",
    "        gen1 = tf.depth_to_space(gen1,2)\n",
    "\n",
    "        gen2 = slim.conv2d(\n",
    "            tf.concat([gen1,enc1],3),num_outputs=128,kernel_size=[3,3],\n",
    "            padding=\"SAME\",normalizer_fn=slim.batch_norm,\n",
    "            activation_fn=tf.nn.elu,weights_initializer=initializer)\n",
    "        gen2 = tf.depth_to_space(gen2,2)\n",
    "\n",
    "        gen3 = slim.conv2d(\n",
    "            tf.concat([gen2,enc0],3),num_outputs=128,kernel_size=[3,3],\n",
    "            padding=\"SAME\",normalizer_fn=slim.batch_norm,\n",
    "            activation_fn=tf.nn.elu, weights_initializer=initializer)\n",
    "        gen3 = tf.depth_to_space(gen3,2)\n",
    "        \n",
    "        g_out = slim.conv2d(\n",
    "            gen3,num_outputs=3,kernel_size=[1,1],padding=\"SAME\",\n",
    "            biases_initializer=None,activation_fn=tf.nn.tanh,\n",
    "            weights_initializer=initializer)\n",
    "        return g_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def discriminator(bottom, reuse=False):\n",
    "    with tf.variable_scope('discriminator'):\n",
    "        filters = [32,64,128,128]\n",
    "        \n",
    "        #Programatically define layers\n",
    "        for i in range(len(filters)):\n",
    "            if i == 0:\n",
    "                layer = slim.conv2d(bottom,filters[i],[3,3],padding=\"SAME\",scope='d'+str(i),\n",
    "                    biases_initializer=None,activation_fn=lrelu,stride=[2,2],\n",
    "                    reuse=reuse,weights_initializer=initializer)\n",
    "            else:\n",
    "                layer = slim.conv2d(bottom,filters[i],[3,3],padding=\"SAME\",scope='d'+str(i),\n",
    "                    normalizer_fn=slim.batch_norm,activation_fn=lrelu,stride=[2,2],\n",
    "                    reuse=reuse,weights_initializer=initializer)\n",
    "            bottom = layer\n",
    "\n",
    "        dis_full = slim.fully_connected(slim.flatten(bottom),1024,activation_fn=lrelu,scope='dl',\n",
    "            reuse=reuse, weights_initializer=initializer)\n",
    "\n",
    "        d_out = slim.fully_connected(dis_full,1,activation_fn=tf.nn.sigmoid,scope='do',\n",
    "            reuse=reuse, weights_initializer=initializer)\n",
    "        return d_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Connecting them together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "#This initializaer is used to initialize all the weights of the network.\n",
    "initializer = tf.truncated_normal_initializer(stddev=0.02)\n",
    "\n",
    "#\n",
    "condition_in = tf.placeholder(shape=[None,height,width,3],dtype=tf.float32)\n",
    "real_in = tf.placeholder(shape=[None,height,width,3],dtype=tf.float32) #Real images\n",
    "\n",
    "Gx = generator(condition_in) #Generates images from random z vectors\n",
    "Dx = discriminator(real_in) #Produces probabilities for real images\n",
    "Dg = discriminator(Gx,reuse=True) #Produces probabilities for generator images\n",
    "\n",
    "#These functions together define the optimization objective of the GAN.\n",
    "d_loss = -tf.reduce_mean(tf.log(Dx) + tf.log(1.-Dg)) #This optimizes the discriminator.\n",
    "#For generator we use traditional GAN objective as well as L1 loss\n",
    "g_loss = -tf.reduce_mean(tf.log(Dg)) + 100*tf.reduce_mean(tf.abs(Gx - real_in)) #This optimizes the generator.\n",
    "\n",
    "#The below code is responsible for applying gradient descent to update the GAN.\n",
    "trainerD = tf.train.AdamOptimizer(learning_rate=0.0002,beta1=0.5)\n",
    "trainerG = tf.train.AdamOptimizer(learning_rate=0.002,beta1=0.5)\n",
    "d_grads = trainerD.compute_gradients(d_loss,slim.get_variables(scope='discriminator'))\n",
    "g_grads = trainerG.compute_gradients(g_loss, slim.get_variables(scope='generator'))\n",
    "\n",
    "update_D = trainerD.apply_gradients(d_grads)\n",
    "update_G = trainerG.apply_gradients(g_grads)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Training the network\n",
    "Now that we have fully defined our network, it is time to train it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "batch_size = 4 #Size of image batch to apply at each iteration.\n",
    "iterations = 500000 #Total number of iterations to use.\n",
    "subset_size = 5000 #How many images to load at a time, will vary depending on available resources\n",
    "frame_directory = './frames' #Directory where training images are located\n",
    "sample_directory = './samples' #Directory to save sample images from generator in.\n",
    "model_directory = './model' #Directory to save trained model to.\n",
    "sample_frequency = 200 #How often to generate sample gif of translated images.\n",
    "save_frequency = 5000 #How often to save model.\n",
    "load_model = False #Whether to load the model or begin training from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "subset = 0\n",
    "dataS = sorted(glob(os.path.join(frame_directory, \"*.png\")))\n",
    "total_subsets = len(dataS)/subset_size\n",
    "init = tf.global_variables_initializer()\n",
    "saver = tf.train.Saver()\n",
    "with tf.Session() as sess:  \n",
    "    sess.run(init)\n",
    "    if load_model == True: \n",
    "        ckpt = tf.train.get_checkpoint_state(model_directory)\n",
    "        saver.restore(sess,ckpt.model_checkpoint_path)\n",
    "\n",
    "    imagesY,imagesX = loadImages(dataS[0:subset_size],False, np.random.randint(0,2)) #Load a subset of images\n",
    "    print \"Loaded subset \" + str(subset)\n",
    "    draw = range(len(imagesX))\n",
    "    for i in range(iterations):\n",
    "        if i % (subset_size/batch_size) != 0 or i == 0:\n",
    "            batch_index = np.random.choice(draw,size=batch_size,replace=False)\n",
    "        else:\n",
    "            subset = np.random.randint(0,total_subsets+1)\n",
    "            imagesY,imagesX = loadImages(dataS[subset*subset_size:(subset+1)*subset_size],False, np.random.randint(0,2))\n",
    "            print \"Loaded subset \" + str(subset)\n",
    "            draw = range(len(imagesX))\n",
    "            batch_index = np.random.choice(draw,size=batch_size,replace=False)\n",
    "        \n",
    "        ys = (np.reshape(imagesY[batch_index],[batch_size,height,width,3]) - 0.5) * 2.0 #Transform to be between -1 and 1\n",
    "        xs = (np.reshape(imagesX[batch_index],[batch_size,height,width,3]) - 0.5) * 2.0\n",
    "        _,dLoss = sess.run([update_D,d_loss],feed_dict={real_in:ys,condition_in:xs}) #Update the discriminator\n",
    "        _,gLoss = sess.run([update_G,g_loss],feed_dict={real_in:ys,condition_in:xs}) #Update the generator\n",
    "        if i % sample_frequency == 0:\n",
    "            print \"Gen Loss: \" + str(gLoss) + \" Disc Loss: \" + str(dLoss)\n",
    "            start_point = np.random.randint(0,len(imagesX)-32)\n",
    "            xs = (np.reshape(imagesX[start_point:start_point+32],[32,height,width,3]) - 0.5) * 2.0\n",
    "            ys = (np.reshape(imagesY[start_point:start_point+32],[32,height,width,3]) - 0.5) * 2.0\n",
    "            sample_G = sess.run(Gx,feed_dict={condition_in:xs}) #Use new z to get sample images from generator.\n",
    "            allS = np.concatenate([xs,sample_G,ys],axis=1)\n",
    "            if not os.path.exists(sample_directory):\n",
    "                os.makedirs(sample_directory)\n",
    "            #Save sample generator images for viewing training progress.\n",
    "            make_gif(allS,'./'+sample_directory+'/a_vid'+str(i)+'.gif',\n",
    "                duration=len(allS)*0.2,true_image=False)\n",
    "        if i % save_frequency == 0 and i != 0:\n",
    "            if not os.path.exists(model_directory):\n",
    "                os.makedirs(model_directory)\n",
    "            saver.save(sess,model_directory+'/model-'+str(i)+'.cptk')\n",
    "            print \"Saved Model\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using a trained network\n",
    "Once we have a trained model saved, we may want to use it to generate new images, and explore the representation it has learned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "test_directory = './test_frames' #Directory to load test frames from\n",
    "subset_size = 5000\n",
    "batch_size = 60 # Size of image batch to apply at each iteration. Will depend of available resources.\n",
    "sample_directory = './test_samples' #Directory to save sample images from generator in.\n",
    "model_directory = './model' #Directory to save trained model to.\n",
    "load_model = True #Whether to load a saved model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dataS = sorted(glob(os.path.join(test_directory, \"*.png\")))\n",
    "subset = 0\n",
    "total_subsets = len(dataS)/subset_size\n",
    "iterations = subset_size / batch_size #Total number of iterations to use.\n",
    "\n",
    "\n",
    "if not os.path.exists(sample_directory):\n",
    "    os.makedirs(sample_directory)\n",
    "\n",
    "init = tf.global_variables_initializer()\n",
    "saver = tf.train.Saver()\n",
    "with tf.Session() as sess:  \n",
    "    sess.run(init)\n",
    "    if load_model == True: \n",
    "        ckpt = tf.train.get_checkpoint_state(model_directory)\n",
    "        saver.restore(sess,ckpt.model_checkpoint_path)\n",
    "        for s in range(total_subsets):\n",
    "            generated_frames = []\n",
    "            _,imagesX = loadImages(dataS[s*subset_size:s*subset_size+subset_size],False, False) #Load a subset of images\n",
    "            for i in range(iterations):\n",
    "                start_point = i * batch_size\n",
    "                xs = (np.reshape(imagesX[start_point:start_point+batch_size],[batch_size,height,width,3]) - 0.5) * 2.0\n",
    "                sample_G = sess.run(Gx,feed_dict={condition_in:xs}) #Use new z to get sample images from generator.    \n",
    "                #allS = np.concatenate([xs,sample_G],axis=2)\n",
    "                generated_frames.append(sample_G)\n",
    "            generated_frames = np.vstack(generated_frames)\n",
    "            for i in range(len(generated_frames)):\n",
    "                im = Image.fromarray(((generated_frames[i]/2.0 + 0.5) * 256).astype('uint8'))\n",
    "                im.save('./'+sample_directory+'/frame'+str(s*subset_size + i)+'.png')  \n",
    "            #make_gif(generated_frames,'./'+sample_directory+'/a_vid'+str(i)+'.gif',\n",
    "            #    duration=len(generated_frames)/10.0,true_image=False)"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}