CycleGAN-keras.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keras implementation of https://junyanz.github.io/CycleGAN/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KERAS_BACKEND']='theano' # 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": {},
   "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": {},
   "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": {},
   "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": {},
   "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": {},
   "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, ndf, max_layers=3, use_sigmoid=True):\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(shape=(nc_in, None, None))\n",
    "    else:\n",
    "        input_a = Input(shape=(None, None, nc_in))\n",
    "    _ = input_a\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\" if use_sigmoid else None) (_)    \n",
    "    return Model(inputs=[input_a], outputs=_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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",
    "use_lsgan = True\n",
    "λ = 10 if use_lsgan else 100\n",
    "\n",
    "loadSize = 143\n",
    "imageSize = 128\n",
    "batchSize = 1\n",
    "lrD = 2e-4\n",
    "lrG = 2e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "netDA = BASIC_D(nc_in, ndf, use_sigmoid = not use_lsgan)\n",
    "netDB = BASIC_D(nc_out, ndf, use_sigmoid = not use_lsgan)\n",
    "netDA.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import SVG\n",
    "from keras.utils.vis_utils import model_to_dot\n",
    "\n",
    "\n",
    "netGB = UNET_G(imageSize, nc_in, nc_out, ngf)\n",
    "netGA = UNET_G(imageSize, nc_out, nc_in, ngf)\n",
    "#SVG(model_to_dot(netG, show_shapes=True).create(prog='dot', format='svg'))\n",
    "netGA.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": {},
   "outputs": [],
   "source": [
    "if use_lsgan:\n",
    "    loss_fn = lambda output, target : K.mean(K.abs(K.square(output-target)))\n",
    "else:\n",
    "    loss_fn = lambda output, target : -K.mean(K.log(output+1e-12)*target+K.log(1-output+1e-12)*(1-target))\n",
    "\n",
    "def cycle_variables(netG1, netG2):\n",
    "    real_input = netG1.inputs[0]\n",
    "    fake_output = netG1.outputs[0]\n",
    "    rec_input = netG2([fake_output])\n",
    "    fn_generate = K.function([real_input], [fake_output, rec_input])\n",
    "    return real_input, fake_output, rec_input, fn_generate\n",
    "\n",
    "real_A, fake_B, rec_A, cycleA_generate = cycle_variables(netGB, netGA)\n",
    "real_B, fake_A, rec_B, cycleB_generate = cycle_variables(netGA, netGB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def D_loss(netD, real, fake, rec):\n",
    "    output_real = netD([real])\n",
    "    output_fake = netD([fake])\n",
    "    loss_D_real = loss_fn(output_real, K.ones_like(output_real))\n",
    "    loss_D_fake = loss_fn(output_fake, K.zeros_like(output_fake))\n",
    "    loss_G = loss_fn(output_fake, K.ones_like(output_fake))\n",
    "    loss_D = loss_D_real+loss_D_fake\n",
    "    loss_cyc = K.mean(K.abs(rec-real))\n",
    "    return loss_D, loss_G, loss_cyc\n",
    "\n",
    "loss_DA, loss_GA, loss_cycA = D_loss(netDA, real_A, fake_A, rec_A)\n",
    "loss_DB, loss_GB, loss_cycB = D_loss(netDB, real_B, fake_B, rec_B)\n",
    "loss_cyc = loss_cycA+loss_cycB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "loss_G = loss_GA+loss_GB+λ*loss_cyc\n",
    "loss_D = loss_DA+loss_DB\n",
    "\n",
    "weightsD = netDA.trainable_weights + netDB.trainable_weights\n",
    "weightsG = netGA.trainable_weights + netGB.trainable_weights\n",
    "\n",
    "training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsD,[],loss_D)\n",
    "netD_train = K.function([real_A, real_B],[loss_DA/2, loss_DB/2], training_updates)\n",
    "training_updates = Adam(lr=lrG, beta_1=0.5).get_updates(weightsG,[], loss_G)\n",
    "netG_train = K.function([real_A, real_B], [loss_GA, loss_GB, loss_cyc], training_updates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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",
    "\n",
    "def read_image(fn):\n",
    "    im = Image.open(fn).convert('RGB')\n",
    "    im = im.resize( (loadSize, 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",
    "    img = arr[h1:h2, w1:w2, :]\n",
    "    if randint(0,1):\n",
    "        img=img[:,::-1]\n",
    "    if channel_first:        \n",
    "        img = np.moveaxis(img, 2, 0)\n",
    "    return img\n",
    "\n",
    "#data = \"edges2shoes\"\n",
    "data = \"horse2zebra\"\n",
    "train_A = load_data('CycleGAN/{}/trainA/*.jpg'.format(data))\n",
    "train_B = load_data('CycleGAN/{}/trainB/*.jpg'.format(data))\n",
    "\n",
    "assert len(train_A) and len(train_B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def minibatch(data, batchsize):\n",
    "    length = len(data)\n",
    "    epoch = i = 0\n",
    "    tmpsize = None    \n",
    "    while True:\n",
    "        size = tmpsize if tmpsize else batchsize\n",
    "        if i+size > length:\n",
    "            shuffle(data)\n",
    "            i = 0\n",
    "            epoch+=1        \n",
    "        rtn = [read_image(data[j]) for j in range(i,i+size)]\n",
    "        i+=size\n",
    "        tmpsize = yield epoch, np.float32(rtn)       \n",
    "\n",
    "def minibatchAB(dataA, dataB, batchsize):\n",
    "    batchA=minibatch(dataA, batchsize)\n",
    "    batchB=minibatch(dataB, batchsize)\n",
    "    tmpsize = None    \n",
    "    while True:        \n",
    "        ep1, A = batchA.send(tmpsize)\n",
    "        ep2, B = batchB.send(tmpsize)\n",
    "        tmpsize = yield max(ep1, ep2), A, B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "train_batch = minibatchAB(train_A, train_B, 6)\n",
    "\n",
    "_, A, B = next(train_batch)\n",
    "showX(A)\n",
    "showX(B)\n",
    "del train_batch, A, B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def showG(A,B):\n",
    "    assert A.shape==B.shape\n",
    "    def G(fn_generate, X):\n",
    "        r = np.array([fn_generate([X[i:i+1]]) for i in range(X.shape[0])])\n",
    "        return r.swapaxes(0,1)[:,:,0]        \n",
    "    rA = G(cycleA_generate, A)\n",
    "    rB = G(cycleB_generate, B)\n",
    "    arr = np.concatenate([A,B,rA[0],rB[0],rA[1],rB[1]])\n",
    "    showX(arr, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from IPython.display import clear_output\n",
    "t0 = time.time()\n",
    "niter = 150\n",
    "gen_iterations = 0\n",
    "epoch = 0\n",
    "errCyc_sum = errGA_sum = errGB_sum = errDA_sum = errDB_sum = 0\n",
    "\n",
    "display_iters = 50\n",
    "#val_batch = minibatch(valAB, 6, direction)\n",
    "train_batch = minibatchAB(train_A, train_B, batchSize)\n",
    "\n",
    "while epoch < niter: \n",
    "    epoch, A, B = next(train_batch)        \n",
    "    errDA, errDB  = netD_train([A, B])\n",
    "    errDA_sum +=errDA\n",
    "    errDB_sum +=errDB\n",
    "\n",
    "    # epoch, trainA, trainB = next(train_batch)\n",
    "    errGA, errGB, errCyc = netG_train([A, B])\n",
    "    errGA_sum += errGA\n",
    "    errGB_sum += errGB\n",
    "    errCyc_sum += errCyc\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 %f Loss_G: %f %f loss_cyc %f'\n",
    "        % (epoch, niter, gen_iterations, errDA_sum/display_iters, errDB_sum/display_iters,\n",
    "           errGA_sum/display_iters, errGB_sum/display_iters, \n",
    "           errCyc_sum/display_iters), time.time()-t0)\n",
    "        _, A, B = train_batch.send(4)\n",
    "        showG(A,B)        \n",
    "        errCyc_sum = errGA_sum = errGB_sum = errDA_sum = errDB_sum = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}