sgan-layerwise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Texture Synthesis with Spatial Generative Adversarial Networks\n",
    "\n",
    "[Paper](https://arxiv.org/pdf/1611.08207v2.pdf)\n",
    "[Sample implementation](https://github.com/ubergmann/spatial_gan)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from nbutil import imshow_multi, to_pil\n",
    "from tensorflow.contrib.layers.python.layers import batch_norm\n",
    "from tensorflow.contrib.layers import xavier_initializer\n",
    "import skimage\n",
    "import skimage.io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 32\n",
    "OUTPUT_WIDTH, OUTPUT_HEIGHT = (256, 256)\n",
    "R = 16\n",
    "Z_WIDTH, Z_HEIGHT = (OUTPUT_WIDTH / R, OUTPUT_HEIGHT / R)\n",
    "Z_DEPTH = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "def load_image(path):\n",
    "    # load image\n",
    "    img = skimage.io.imread(path)\n",
    "    img = img[:,:,:3] # drop alpha channel\n",
    "    return img / 255.0\n",
    "\n",
    "image = load_image('../data/bk.png')\n",
    "imshow_multi([image])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def _phase_shift(I, r):\n",
    "    # Helper function with main phase shift operation\n",
    "    bsize, a, b, c = I.get_shape().as_list()\n",
    "    X = tf.reshape(I, (bsize, a, b, r, r))\n",
    "    X = tf.transpose(X, (0, 1, 2, 4, 3))  # bsize, a, b, 1, 1\n",
    "    X = tf.split(1, a, X)  # a, [bsize, b, r, r]\n",
    "    X = tf.concat(2, [tf.squeeze(x) for x in X])  # bsize, b, a*r, r\n",
    "    X = tf.split(1, b, X)  # b, [bsize, a*r, r]\n",
    "    X = tf.concat(2, [tf.squeeze(x) for x in X])  #\n",
    "    bsize, a*r, b*r\n",
    "    return tf.reshape(X, (bsize, a*r, b*r, 1))\n",
    "\n",
    "def PS(X, r, color=False):\n",
    "    # Main OP that you can arbitrarily use in you tensorflow code\n",
    "    if color:\n",
    "        Xc = tf.split(3, 3, X)\n",
    "        X = tf.concat(3, [_phase_shift(x, r) for x in Xc])\n",
    "    else:\n",
    "        X = _phase_shift(X, r)\n",
    "    return X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dropout_keep_prob = tf.placeholder_with_default(tf.constant(1.0), [], name='dropout_keep_prob')        \n",
    "\n",
    "def lrelu(x):\n",
    "    alpha = 0.05\n",
    "    return tf.maximum(alpha*x, x)\n",
    "\n",
    "def create_batch_norm(inputs, name='bn'):\n",
    "    with tf.variable_scope(name):\n",
    "        return batch_norm(inputs, is_training=True, updates_collections=None)\n",
    "\n",
    "def create_avg_pool(inputs, ksize=2, stride=2):\n",
    "    return tf.nn.avg_pool(inputs, ksize=[1, ksize, ksize, 1], strides=[1, stride, stride, 1], padding='SAME')\n",
    "    \n",
    "def create_dropout(inputs):\n",
    "    return tf.nn.dropout(inputs, dropout_keep_prob)\n",
    "\n",
    "def create_conv(input, out_channels, patch_size=5, stride=1, use_relu=True, name='conv'):\n",
    "    with tf.variable_scope(name):\n",
    "        in_channels = input.get_shape()[-1].value\n",
    "        # w = weight_var([patch_size, patch_size, in_channels, out_channels], name='w', key=join_keys(key, 'w'))\n",
    "        # b = weight_var([out_channels], stddev=0, name='b', mean=0.1, key=join_keys(key, 'b'))\n",
    "        w = tf.get_variable('w', \n",
    "                            shape=[patch_size, patch_size, in_channels, out_channels], \n",
    "                            initializer=xavier_initializer())\n",
    "        b = tf.get_variable('b',\n",
    "                           shape=[out_channels])\n",
    "        conv = tf.nn.conv2d(input, w, strides=[1,stride,stride,1], padding='SAME')\n",
    "        activation = lrelu(conv + b) if use_relu else conv + b\n",
    "        return activation\n",
    "\n",
    "def create_deconv(input, out_channels, patch_size=5, stride=1, use_relu=True, name='deconv'):\n",
    "    with tf.variable_scope(name):\n",
    "        # for best results, patch_size should be a multiple of stride\n",
    "        input_w, input_h, input_channels = [i.value for i in input.get_shape()[-3:]]\n",
    "\n",
    "        # w = weight_var([patch_size, patch_size, out_channels, input_channels])\n",
    "        # b = weight_var([out_channels], mean=0.1)\n",
    "        w = tf.get_variable('w', \n",
    "                            shape=[patch_size, patch_size, out_channels, input_channels], \n",
    "                            initializer=xavier_initializer())\n",
    "        b = tf.get_variable('b',\n",
    "                            shape=[out_channels])\n",
    "\n",
    "        batch_size = BATCH_SIZE # tf.shape(input)[0]\n",
    "        output_shape = tf.pack([batch_size, input_w*stride, input_h*stride, out_channels])\n",
    "\n",
    "        deconv = tf.nn.conv2d_transpose(input, w, output_shape, strides=[1,stride,stride,1], padding='SAME')\n",
    "\n",
    "        activation = lrelu(deconv + b) if use_relu else deconv + b\n",
    "        return activation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "source_image = tf.constant(image, tf.float32)\n",
    "def rand_crop(): return tf.random_crop(source_image, [OUTPUT_WIDTH, OUTPUT_HEIGHT, 3])\n",
    "real_textures = [rand_crop() for _ in xrange(BATCH_SIZE)]\n",
    "\n",
    "def generator(noise):\n",
    "    with tf.variable_scope('generator'):\n",
    "        image = noise\n",
    "        # strides must multiply to R (16)\n",
    "        layers = [256, 128, 64, 3]\n",
    "        for i, channels in enumerate(layers):\n",
    "            is_last_layer = i == len(layers)-1\n",
    "            image = create_deconv(image, \n",
    "                                  channels, \n",
    "                                  size=5, \n",
    "                                  stride=2, \n",
    "                                  name='deconv'+str(i), \n",
    "                                  use_relu=(not is_last_layer))\n",
    "            if not is_last_layer:\n",
    "                image = create_batch_norm(image, name='bn'+str(i))\n",
    "        assert [d.value for d in image.get_shape()[1:]] == [OUTPUT_WIDTH, OUTPUT_HEIGHT, 3]\n",
    "        return image\n",
    "\n",
    "def discriminator(textures):\n",
    "    # input: [BATCH_SIZE, IMAGE_WIDTH, IMAGE_HEIGHT, 3] textures\n",
    "    # output: [BATCH_SIZE, Z_WIDTH, Z_HEIGHT, 2]\n",
    "    #  where the last dim is the logit probability of this texture being real\n",
    "    with tf.variable_scope('discriminator'):\n",
    "        image = textures\n",
    "        # some big convolutional layers:\n",
    "        for i, channels in enumerate([256, 128, 64, 2]):\n",
    "            image = create_conv(image, channels, stride=1, patch_size=5, name='conv'+str(i))\n",
    "            image = create_avg_pool(image)\n",
    "            image = create_batch_norm(image, name='bn'+str(i))   \n",
    "        # a couple 1x1 convolutions:\n",
    "        for i, channels in enumerate([2]):\n",
    "            image = create_conv(image, channels, stride=1, patch_size=1, name='1x1conv'+str(i))\n",
    "        assert [d.value for d in image.get_shape()[1:]] == [Z_WIDTH, Z_HEIGHT, 2]\n",
    "        return image\n",
    "\n",
    "scopename = '3'\n",
    "with tf.variable_scope(scopename):\n",
    "\n",
    "    noise = tf.random_uniform([BATCH_SIZE, Z_WIDTH, Z_HEIGHT, Z_DEPTH], minval=-1, maxval=1)\n",
    "    synthetic_textures = generator(noise)\n",
    "\n",
    "    disc_input = tf.concat(0, [real_textures, synthetic_textures])\n",
    "    disc_target = tf.concat(0, [tf.ones([BATCH_SIZE, Z_WIDTH, Z_HEIGHT], tf.int32), tf.zeros([BATCH_SIZE, Z_WIDTH, Z_HEIGHT], tf.int32)])\n",
    "    \n",
    "    disc_output = discriminator(disc_input)\n",
    "    discriminator_loss = tf.reduce_sum(tf.nn.sparse_softmax_cross_entropy_with_logits(disc_output, disc_target))\n",
    "\n",
    "    disc_guess = tf.argmax(disc_output, 3)\n",
    "    disc_correct = tf.equal(tf.cast(disc_guess, tf.int32), disc_target)\n",
    "    disc_accuracy = tf.reduce_mean(tf.cast(disc_correct, tf.float32))\n",
    "\n",
    "    disc_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scopename+'/discriminator')\n",
    "    gen_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scopename+'/generator')\n",
    "    # for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):\n",
    "    #     print v.name\n",
    "\n",
    "    global_step = tf.contrib.framework.get_or_create_global_step()\n",
    "\n",
    "    train_gen = tf.train.AdamOptimizer(0.0001).minimize(-discriminator_loss, global_step=global_step, var_list=gen_vars)\n",
    "    train_disc = tf.train.AdamOptimizer(0.0001).minimize(discriminator_loss, global_step=global_step, var_list=disc_vars)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "session = tf.InteractiveSession()\n",
    "\n",
    "save_path = 'models/sgan256-2-1'\n",
    "\n",
    "init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())\n",
    "session.run(init_op)\n",
    "\n",
    "import os\n",
    "saver = None\n",
    "if save_path:\n",
    "    if not os.path.exists(save_path):\n",
    "        os.mkdir(save_path)\n",
    "    saver = tf.train.Saver()\n",
    "    ckpt = tf.train.get_checkpoint_state(save_path)\n",
    "    if ckpt and ckpt.model_checkpoint_path:\n",
    "        saver.restore(session, ckpt.model_checkpoint_path)\n",
    "        print 'Restored from checkpoint', ckpt.model_checkpoint_path\n",
    "    else:\n",
    "        print 'Did not restore from checkpoint'\n",
    "else:\n",
    "    print 'Will not save progress'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def avg(x): return sum(x) / float(len(x))\n",
    "disc_accuracies = []\n",
    "losses = []\n",
    "last_saved_loss = None\n",
    "\n",
    "while True:\n",
    "    only_train_gen = len(disc_accuracies) > 0 and avg(disc_accuracies) > 0.8\n",
    "    only_train_disc = len(disc_accuracies) > 0 and avg(disc_accuracies) < 0.6\n",
    "    n_gen_runs = 0 if only_train_disc else 1\n",
    "    n_disc_runs = 0 if only_train_gen else 1\n",
    "        \n",
    "    for _ in range(n_gen_runs):\n",
    "        disc_acc_, step_, loss_, _ = session.run([disc_accuracy, global_step, discriminator_loss, train_gen])\n",
    "        losses.append(loss_)\n",
    "        disc_accuracies.append(disc_acc_)\n",
    "    \n",
    "    for _ in range(n_disc_runs):\n",
    "        feed = {dropout_keep_prob: 0.5}\n",
    "        disc_acc_, step_, loss_, _ = session.run([disc_accuracy, global_step, discriminator_loss, train_disc], feed_dict=feed)\n",
    "        losses.append(loss_)\n",
    "        disc_accuracies.append(disc_acc_)\n",
    "    \n",
    "    step_rounded = int(step_ / 2) * 2\n",
    "    if step_rounded % 50 == 0:\n",
    "        print \"Step: {}, loss: {}, disc accuracy: {}\".format(step_rounded, avg(losses), avg(disc_accuracies))\n",
    "        \n",
    "        if step_rounded % 200 == 0 and saver:\n",
    "            should_save = True\n",
    "            if should_save:\n",
    "                saver.save(session, save_path + '/model.ckpt', global_step=step_rounded)\n",
    "                print 'Saved'\n",
    "            else:\n",
    "                pass\n",
    "                # print 'Loss did not decrease from previous save, so not saving'\n",
    "        \n",
    "        disc_accuracies = []\n",
    "        losses = []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "def generate_sample():\n",
    "    textures, real = session.run([synthetic_textures, real_textures])\n",
    "    imshow_multi(list(np.clip(textures[:3], 0, 1)) + list(real[:1]))\n",
    "    return to_pil(textures[0])\n",
    "generate_sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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": 2
}