pix2pix-keras.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keras implementation of https://phillipi.github.io/pix2pix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KERAS_BACKEND']='tensorflow' # can choose theano, tensorflow, cntk\n",
    "os.environ['THEANO_FLAGS']='floatX=float32,device=cuda,optimizer=fast_run,dnn.library_path=/usr/lib'\n",
    "#os.environ['THEANO_FLAGS']='floatX=float32,device=cuda,optimizer=fast_compile,dnn.library_path=/usr/lib'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "if os.environ['KERAS_BACKEND'] =='theano':\n",
    "    channel_axis=1\n",
    "    K.set_image_data_format('channels_first')\n",
    "    channel_first = True\n",
    "else:\n",
    "    K.set_image_data_format('channels_last')\n",
    "    channel_axis=-1\n",
    "    channel_first = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from keras.models import Sequential, Model\n",
    "from keras.layers import Conv2D, ZeroPadding2D, BatchNormalization, Input, Dropout\n",
    "from keras.layers import Conv2DTranspose, Reshape, Activation, Cropping2D, Flatten\n",
    "from keras.layers import Concatenate\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "from keras.activations import relu\n",
    "from keras.initializers import RandomNormal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Weights initializations\n",
    "# bias are initailized as 0\n",
    "def __conv_init(a):\n",
    "    print(\"conv_init\", a)\n",
    "    k = RandomNormal(0, 0.02)(a) # for convolution kernel\n",
    "    k.conv_weight = True    \n",
    "    return k\n",
    "conv_init = RandomNormal(0, 0.02)\n",
    "gamma_init = RandomNormal(1., 0.02) # for batch normalization\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# HACK speed up theano\n",
    "if K._BACKEND == 'theano':\n",
    "    import keras.backend.theano_backend as theano_backend\n",
    "    def _preprocess_conv2d_kernel(kernel, data_format):\n",
    "        #return kernel\n",
    "        if hasattr(kernel, \"original\"):\n",
    "            print(\"use original\")\n",
    "            return kernel.original\n",
    "        elif hasattr(kernel, '_keras_shape'):\n",
    "            s = kernel._keras_shape\n",
    "            print(\"use reshape\",s)\n",
    "            kernel = kernel.reshape((s[3], s[2],s[0], s[1]))\n",
    "        else:\n",
    "            kernel = kernel.dimshuffle((3, 2, 0, 1))\n",
    "        return kernel\n",
    "    theano_backend._preprocess_conv2d_kernel = _preprocess_conv2d_kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Basic discriminator\n",
    "def conv2d(f, *a, **k):\n",
    "    return Conv2D(f, kernel_initializer = conv_init, *a, **k)\n",
    "def batchnorm():\n",
    "    return BatchNormalization(momentum=0.9, axis=channel_axis, epsilon=1.01e-5,\n",
    "                                   gamma_initializer = gamma_init)\n",
    "def BASIC_D(nc_in, nc_out, ndf, max_layers=3):\n",
    "    \"\"\"DCGAN_D(nc, ndf, max_layers=3)\n",
    "       nc: channels\n",
    "       ndf: filters of the first layer\n",
    "       max_layers: max hidden layers\n",
    "    \"\"\"    \n",
    "    if channel_first:\n",
    "        input_a, input_b =  Input(shape=(nc_in, None, None)), Input(shape=(nc_out, None, None))\n",
    "    else:\n",
    "        input_a, input_b = Input(shape=(None, None, nc_in)), Input(shape=(None, None, nc_out))\n",
    "    _ = Concatenate(axis=channel_axis)([input_a, input_b])\n",
    "    _ = conv2d(ndf, kernel_size=4, strides=2, padding=\"same\", name = 'First') (_)\n",
    "    _ = LeakyReLU(alpha=0.2)(_)\n",
    "    \n",
    "    for layer in range(1, max_layers):        \n",
    "        out_feat = ndf * min(2**layer, 8)\n",
    "        _ = conv2d(out_feat, kernel_size=4, strides=2, padding=\"same\", \n",
    "                   use_bias=False, name = 'pyramid.{0}'.format(layer)             \n",
    "                        ) (_)\n",
    "        _ = batchnorm()(_, training=1)        \n",
    "        _ = LeakyReLU(alpha=0.2)(_)\n",
    "    \n",
    "    out_feat = ndf*min(2**max_layers, 8)\n",
    "    _ = ZeroPadding2D(1)(_)\n",
    "    _ = conv2d(out_feat, kernel_size=4,  use_bias=False, name = 'pyramid_last') (_)\n",
    "    _ = batchnorm()(_, training=1)\n",
    "    _ = LeakyReLU(alpha=0.2)(_)\n",
    "    \n",
    "    # final layer\n",
    "    _ = ZeroPadding2D(1)(_)\n",
    "    _ = conv2d(1, kernel_size=4, name = 'final'.format(out_feat, 1), \n",
    "               activation = \"sigmoid\") (_)    \n",
    "    return Model(inputs=[input_a, input_b], outputs=_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def UNET_G(isize, nc_in=3, nc_out=3, ngf=64, fixed_input_size=True):    \n",
    "    max_nf = 8*ngf    \n",
    "    def block(x, s, nf_in, use_batchnorm=True, nf_out=None, nf_next=None):\n",
    "        # print(\"block\",x,s,nf_in, use_batchnorm, nf_out, nf_next)\n",
    "        assert s>=2 and s%2==0\n",
    "        if nf_next is None:\n",
    "            nf_next = min(nf_in*2, max_nf)\n",
    "        if nf_out is None:\n",
    "            nf_out = nf_in\n",
    "        x = conv2d(nf_next, kernel_size=4, strides=2, use_bias=(not (use_batchnorm and s>2)),\n",
    "                   padding=\"same\", name = 'conv_{0}'.format(s)) (x)\n",
    "        if s>2:\n",
    "            if use_batchnorm:\n",
    "                x = batchnorm()(x, training=1)\n",
    "            x2 = LeakyReLU(alpha=0.2)(x)\n",
    "            x2 = block(x2, s//2, nf_next)\n",
    "            x = Concatenate(axis=channel_axis)([x, x2])            \n",
    "        x = Activation(\"relu\")(x)\n",
    "        x = Conv2DTranspose(nf_out, kernel_size=4, strides=2, use_bias=not use_batchnorm,\n",
    "                            kernel_initializer = conv_init,          \n",
    "                            name = 'convt.{0}'.format(s))(x)        \n",
    "        x = Cropping2D(1)(x)\n",
    "        if use_batchnorm:\n",
    "            x = batchnorm()(x, training=1)\n",
    "        if s <=8:\n",
    "            x = Dropout(0.5)(x, training=1)\n",
    "        return x\n",
    "    \n",
    "    s = isize if fixed_input_size else None\n",
    "    if channel_first:\n",
    "        _ = inputs = Input(shape=(nc_in, s, s))\n",
    "    else:\n",
    "        _ = inputs = Input(shape=(s, s, nc_in))        \n",
    "    _ = block(_, isize, nc_in, False, nf_out=nc_out, nf_next=ngf)\n",
    "    _ = Activation('tanh')(_)\n",
    "    return Model(inputs=inputs, outputs=[_])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "nc_in = 3\n",
    "nc_out = 3\n",
    "ngf = 64\n",
    "ndf = 64\n",
    "λ = 10\n",
    "\n",
    "loadSize = 286\n",
    "imageSize = 256\n",
    "batchSize = 1\n",
    "lrD = 2e-4\n",
    "lrG = 2e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "netD = BASIC_D(nc_in, nc_out, ndf)\n",
    "netD.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import SVG\n",
    "from keras.utils.vis_utils import model_to_dot\n",
    "\n",
    "\n",
    "netG = UNET_G(imageSize, nc_in, nc_out, ngf)\n",
    "#SVG(model_to_dot(netG, show_shapes=True).create(prog='dot', format='svg'))\n",
    "netG.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from keras.optimizers import RMSprop, SGD, Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "real_A = netG.input\n",
    "fake_B = netG.output\n",
    "netG_generate = K.function([real_A], [fake_B])\n",
    "real_B = netD.inputs[1]\n",
    "output_D_real = netD([real_A, real_B])\n",
    "output_D_fake = netD([real_A, fake_B])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#loss_fn = lambda output, target : K.mean(K.binary_crossentropy(output, target))\n",
    "loss_fn = lambda output, target : -K.mean(K.log(output+1e-12)*target+K.log(1-output+1e-12)*(1-target))\n",
    "\n",
    "loss_D_real = loss_fn(output_D_real, K.ones_like(output_D_real))\n",
    "loss_D_fake = loss_fn(output_D_fake, K.zeros_like(output_D_fake))\n",
    "loss_G_fake = loss_fn(output_D_fake, K.ones_like(output_D_fake))\n",
    "\n",
    "\n",
    "loss_L1 = K.mean(K.abs(fake_B-real_B))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loss_D = loss_D_real +loss_D_fake\n",
    "training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(netD.trainable_weights,[],loss_D)\n",
    "netD_train = K.function([real_A, real_B],[loss_D/2], training_updates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loss_G = loss_G_fake   + 100 * loss_L1\n",
    "training_updates = Adam(lr=lrG, beta_1=0.5).get_updates(netG.trainable_weights,[], loss_G)\n",
    "netG_train = K.function([real_A, real_B], [loss_G_fake, loss_L1], training_updates)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "import glob\n",
    "from random import randint, shuffle\n",
    "\n",
    "def load_data(file_pattern):\n",
    "    return glob.glob(file_pattern)\n",
    "def read_image(fn, direction=0):\n",
    "    im = Image.open(fn)\n",
    "    im = im.resize( (loadSize*2, loadSize), Image.BILINEAR )\n",
    "    arr = np.array(im)/255*2-1\n",
    "    w1,w2 = (loadSize-imageSize)//2,(loadSize+imageSize)//2\n",
    "    h1,h2 = w1,w2\n",
    "    imgA = arr[h1:h2, loadSize+w1:loadSize+w2, :]\n",
    "    imgB = arr[h1:h2, w1:w2, :]\n",
    "    if randint(0,1):\n",
    "        imgA=imgA[:,::-1]\n",
    "        imgB=imgB[:,::-1]\n",
    "    if channel_first:\n",
    "        imgA = np.moveaxis(imgA, 2, 0)\n",
    "        imgB = np.moveaxis(imgB, 2, 0)\n",
    "    if direction==0:\n",
    "        return imgA, imgB\n",
    "    else:\n",
    "        return imgB,imgA\n",
    "\n",
    "data = \"edges2shoes\"\n",
    "data = \"facades\"\n",
    "direction = 0\n",
    "trainAB = load_data('pix2pix/{}/train/*.jpg'.format(data))\n",
    "valAB = load_data('pix2pix/{}/val/*.jpg'.format(data))\n",
    "assert len(trainAB) and len(valAB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def minibatch(dataAB, batchsize, direction=0):\n",
    "    length = len(dataAB)\n",
    "    epoch = i = 0\n",
    "    tmpsize = None    \n",
    "    while True:\n",
    "        size = tmpsize if tmpsize else batchsize\n",
    "        if i+size > length:\n",
    "            shuffle(dataAB)\n",
    "            i = 0\n",
    "            epoch+=1        \n",
    "        dataA = []\n",
    "        dataB = []\n",
    "        for j in range(i,i+size):\n",
    "            imgA,imgB = read_image(dataAB[j], direction)\n",
    "            dataA.append(imgA)\n",
    "            dataB.append(imgB)\n",
    "        dataA = np.float32(dataA)\n",
    "        dataB = np.float32(dataB)\n",
    "        i+=size\n",
    "        tmpsize = yield epoch, dataA, dataB        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "def showX(X, rows=1):\n",
    "    assert X.shape[0]%rows == 0\n",
    "    int_X = ( (X+1)/2*255).clip(0,255).astype('uint8')\n",
    "    if channel_first:\n",
    "        int_X = np.moveaxis(int_X.reshape(-1,3,imageSize,imageSize), 1, 3)\n",
    "    else:\n",
    "        int_X = int_X.reshape(-1,imageSize,imageSize, 3)\n",
    "    int_X = int_X.reshape(rows, -1, imageSize, imageSize,3).swapaxes(1,2).reshape(rows*imageSize,-1, 3)\n",
    "    display(Image.fromarray(int_X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "train_batch = minibatch(trainAB, 6, direction=direction)\n",
    "_, trainA, trainB = next(train_batch)\n",
    "showX(trainA)\n",
    "showX(trainB)\n",
    "del train_batch, trainA, trainB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def netG_gen(A):\n",
    "    return np.concatenate([netG_generate([A[i:i+1]])[0] for i in range(A.shape[0])], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from IPython.display import clear_output\n",
    "t0 = time.time()\n",
    "niter = 50\n",
    "gen_iterations = 0\n",
    "errL1 = epoch = errG = 0\n",
    "errL1_sum = errG_sum = errD_sum = 0\n",
    "\n",
    "display_iters = 500\n",
    "val_batch = minibatch(valAB, 6, direction)\n",
    "train_batch = minibatch(trainAB, batchSize, direction)\n",
    "\n",
    "while epoch < niter: \n",
    "    epoch, trainA, trainB = next(train_batch)        \n",
    "    errD,  = netD_train([trainA, trainB])\n",
    "    errD_sum +=errD\n",
    "\n",
    "    errG, errL1 = netG_train([trainA, trainB])\n",
    "    errG_sum += errG\n",
    "    errL1_sum += errL1\n",
    "    gen_iterations+=1\n",
    "    if gen_iterations%display_iters==0:\n",
    "        if gen_iterations%(5*display_iters)==0:\n",
    "            clear_output()\n",
    "        print('[%d/%d][%d] Loss_D: %f Loss_G: %f loss_L1: %f'\n",
    "        % (epoch, niter, gen_iterations, errD_sum/display_iters, errG_sum/display_iters, errL1_sum/display_iters), time.time()-t0)\n",
    "        _, valA, valB = train_batch.send(6) \n",
    "        fakeB = netG_gen(valA)\n",
    "        showX(np.concatenate([valA, valB, fakeB], axis=0), 3)\n",
    "        errL1_sum = errG_sum = errD_sum = 0\n",
    "        _, valA, valB = next(val_batch)\n",
    "        fakeB = netG_gen(valA)\n",
    "        showX(np.concatenate([valA, valB, fakeB], axis=0), 3)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "_, valA, valB = train_batch.send(6) \n",
    "fakeB = netG_gen(valA)\n",
    "showX(np.concatenate([valA, valB, fakeB], axis=0), 3)\n",
    "errL1_sum = errG_sum = errD_sum = 0\n",
    "_, valA, valB = next(val_batch)\n",
    "fakeB = netG_gen(valA)\n",
    "showX(np.concatenate([valA, valB, fakeB], axis=0), 3)"
   ]
  }
 ],
 "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}