{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "from glob import glob\n",
    "import numpy as np\n",
    "from matplotlib import pyplot\n",
    "from PIL import Image\n",
    "import tensorflow as tf\n",
    "\n",
    "##README : IF output folder already existed in same route, it makes error. change past output folder's name ##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset(object):\n",
    "    def __init__(self, data_files):\n",
    "        IMAGE_WIDTH = 25\n",
    "        IMAGE_HEIGHT = 25\n",
    "        self.image_mode = 'RGB'\n",
    "        image_channels = 3\n",
    "        self.data_files = data_files\n",
    "        self.shape = len(data_files), IMAGE_WIDTH, IMAGE_HEIGHT, image_channels\n",
    "\n",
    "    def get_batches(self, batch_size):\n",
    "        IMAGE_MAX_VALUE = 255\n",
    "        current_index = 0\n",
    "        while current_index + batch_size <= self.shape[0]:\n",
    "            data_batch = get_batch(\n",
    "                self.data_files[current_index:current_index + batch_size],\n",
    "                self.shape[1],self.shape[2],\n",
    "                self.image_mode)\n",
    "            \n",
    "            current_index += batch_size\n",
    "            \n",
    "            yield data_batch / IMAGE_MAX_VALUE - 0.5\n",
    "\n",
    "\n",
    "def model_inputs(image_width, image_height, image_channels, z_dim):\n",
    "    real_input_images = tf.placeholder(tf.float32, [None, image_width, image_height, image_channels], 'real_input_images')\n",
    "    input_z = tf.placeholder(tf.float32, [None, z_dim], 'input_z')\n",
    "    learning_rate = tf.placeholder(tf.float32, [], 'learning_rate')\n",
    "    return real_input_images, input_z, learning_rate\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator(images, reuse=False, alpha=0.2, keep_prob=0.5):\n",
    "    with tf.variable_scope('discriminator', reuse=reuse):\n",
    "        # Input layer is 25x25xn\n",
    "        # Convolutional layer, 13x13x64\n",
    "        conv1 = tf.layers.conv2d(images, 64, 5, 2, padding='same', kernel_initializer=tf.contrib.layers.xavier_initializer())\n",
    "        lrelu1 = tf.maximum(alpha * conv1, conv1)\n",
    "        drop1 = tf.layers.dropout(lrelu1, keep_prob)\n",
    "        \n",
    "        # Strided convolutional layer, 7x7x128\n",
    "        conv2 = tf.layers.conv2d(drop1, 128, 5, 2, 'same', use_bias=False)\n",
    "        bn2 = tf.layers.batch_normalization(conv2)\n",
    "        lrelu2 = tf.maximum(alpha * bn2, bn2)\n",
    "        drop2 = tf.layers.dropout(lrelu2, keep_prob)\n",
    "        \n",
    "        # Strided convolutional layer, 4x4x256\n",
    "        conv3 = tf.layers.conv2d(drop2, 256, 5, 2, 'same', use_bias=False)\n",
    "        bn3 = tf.layers.batch_normalization(conv3)\n",
    "        lrelu3 = tf.maximum(alpha * bn3, bn3)\n",
    "        drop3 = tf.layers.dropout(lrelu3, keep_prob)\n",
    "        \n",
    "        # fully connected\n",
    "        flat = tf.reshape(drop3, (-1, 4*4*256))\n",
    "        logits = tf.layers.dense(flat, 1)\n",
    "        out = tf.sigmoid(logits)\n",
    "        \n",
    "        return out, logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator(z, out_channel_dim, is_train=True, alpha=0.2, keep_prob=0.5):\n",
    "    # TODO: Implement Function\n",
    "    with tf.variable_scope('generator', reuse=(not is_train)):\n",
    "        # First fully connected layer, 8x4x512\n",
    "        fc = tf.layers.dense(z, 4*4*1024, use_bias=False)\n",
    "        fc = tf.reshape(fc, (-1, 4, 4, 1024))\n",
    "        bn0 = tf.layers.batch_normalization(fc, training=is_train)\n",
    "        lrelu0 = tf.maximum(alpha * bn0, bn0)\n",
    "        drop0 = tf.layers.dropout(lrelu0, keep_prob, training=is_train)\n",
    "        \n",
    "        # Deconvolution, 16x8x256\n",
    "        conv1 = tf.layers.conv2d_transpose(drop0, 512,3, 1, 'valid', use_bias=False)\n",
    "        bn1 = tf.layers.batch_normalization(conv1, training=is_train)\n",
    "        lrelu1 = tf.maximum(alpha * bn1, bn1)\n",
    "        drop1 = tf.layers.dropout(lrelu1, keep_prob, training=is_train)\n",
    "        \n",
    "        # Deconvolution, 32x 128\n",
    "        conv2 = tf.layers.conv2d_transpose(drop1, 256, 3, 2, 'same', use_bias=False)\n",
    "        bn2 = tf.layers.batch_normalization(conv2, training=is_train)\n",
    "        lrelu2 = tf.maximum(alpha * bn2, bn2)\n",
    "        drop2 = tf.layers.dropout(lrelu2, keep_prob, training=is_train)\n",
    "        \n",
    "        # Output layer, 28x28xn\n",
    "        logits = tf.layers.conv2d_transpose(drop2, out_channel_dim, 3, 2, 'valid')\n",
    "        \n",
    "        out = tf.tanh(logits)\n",
    "        \n",
    "        print(fc.shape)\n",
    "        print(drop1.shape)\n",
    "        print(drop2.shape)\n",
    "        print(logits.shape)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_loss(input_real, input_z, out_channel_dim, alpha=0.2, smooth_factor=0.1):\n",
    "    d_model_real, d_logits_real = discriminator(input_real, alpha=alpha)\n",
    "    \n",
    "    d_loss_real = tf.reduce_mean(\n",
    "        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real,\n",
    "                                                labels=tf.ones_like(d_model_real) * (1 - smooth_factor)))\n",
    "    \n",
    "    input_fake = generator(input_z, out_channel_dim, alpha=alpha)\n",
    "    d_model_fake, d_logits_fake = discriminator(input_fake, reuse=True, alpha=alpha)\n",
    "    \n",
    "    d_loss_fake = tf.reduce_mean(\n",
    "        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.zeros_like(d_model_fake)))\n",
    "    \n",
    "    g_loss = tf.reduce_mean(\n",
    "        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_model_fake)))\n",
    "\n",
    "    return d_loss_real + d_loss_fake, g_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_opt(d_loss, g_loss, learning_rate, beta1):\n",
    "    # Get weights and bias to update\n",
    "    t_vars = tf.trainable_variables()\n",
    "    d_vars = [var for var in t_vars if var.name.startswith('discriminator')]\n",
    "    g_vars = [var for var in t_vars if var.name.startswith('generator')]\n",
    "\n",
    "    # Optimize\n",
    "    with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):\n",
    "        d_train_opt = tf.train.AdamOptimizer(learning_rate, beta1=beta1).minimize(d_loss, var_list=d_vars)\n",
    "        g_train_opt = tf.train.AdamOptimizer(learning_rate, beta1=beta1).minimize(g_loss, var_list=g_vars)\n",
    "\n",
    "    return d_train_opt, g_train_opt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_generator_output(sess, n_images, input_z, out_channel_dim, image_mode):\n",
    "    cmap = None if image_mode == 'RGB' else 'gray'\n",
    "    z_dim = input_z.get_shape().as_list()[-1]\n",
    "    example_z = np.random.uniform(-1, 1, size=[n_images, z_dim])\n",
    "\n",
    "    samples = sess.run(\n",
    "        generator(input_z, out_channel_dim, False),\n",
    "        feed_dict={input_z: example_z})\n",
    "    \n",
    "    # pyplot.show()\n",
    "    return samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch_count, batch_size, z_dim, learning_rate, beta1, get_batches, data_shape, data_image_mode,\n",
    "          print_every=10, show_every=10):\n",
    "    # TODO: Build Model\n",
    "    input_real, input_z, _ = model_inputs(data_shape[2], data_shape[1], data_shape[3], z_dim)\n",
    "    d_loss, g_loss = model_loss(input_real, input_z, data_shape[3], alpha=0.2)\n",
    "    d_train_opt, g_train_opt = model_opt(d_loss, g_loss, learning_rate, beta1)\n",
    "    \n",
    "    saver = tf.train.Saver()\n",
    "    sample_z = np.random.uniform(-1, 1, size=(72, z_dim))\n",
    "    \n",
    "    samples, losses = [], []\n",
    "    \n",
    "    steps = 0\n",
    "    count = 0\n",
    "    \n",
    "    with tf.Session() as sess:\n",
    "        saver = tf.train.Saver()\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        \n",
    "         # continue training\n",
    "        save_path = saver.save(sess, \"/tmp/model.ckpt\")\n",
    "        ckpt = tf.train.latest_checkpoint('./model/')\n",
    "        saver.restore(sess, save_path)\n",
    "        coord = tf.train.Coordinator()\n",
    "        threads = tf.train.start_queue_runners(sess=sess, coord=coord)\n",
    "\n",
    "        os.mkdir('output')\n",
    "        for epoch_i in range(epoch_count):\n",
    "            for batch_images in get_batches(batch_size):\n",
    "                # Train Model\n",
    "                steps += 1\n",
    "                batch_images *= 2.0\n",
    "                \n",
    "                # Sample random noise for G\n",
    "                batch_z = np.random.uniform(-1, 1, size=(batch_size, z_dim))\n",
    "                \n",
    "                # Run optimizers\n",
    "                sess.run(d_train_opt, feed_dict={input_real: batch_images, input_z: batch_z})\n",
    "                sess.run(g_train_opt, feed_dict={input_z: batch_z})\n",
    "                \n",
    "                if steps % print_every == 0:\n",
    "                    os.mkdir('output/'+ str(steps))\n",
    "                    # At the end of each epoch, get the losses and print them out\n",
    "                    train_loss_d = d_loss.eval({input_real: batch_images, input_z: batch_z})\n",
    "                    train_loss_g = g_loss.eval({input_z: batch_z})\n",
    "                    print(\"Epoch {}/{} Step {}...\".format(epoch_i+1, epoch_count, steps),\n",
    "                      \"Discriminator Loss: {:.4f}...\".format(train_loss_d),\n",
    "                      \"Generator Loss: {:.4f}\".format(train_loss_g))\n",
    "                    # Save losses for viewing after training\n",
    "                    #losses.append((train_loss_d, train_loss_g))\n",
    "                            \n",
    "                if steps % show_every == 0:\n",
    "                    count = count +1\n",
    "                    iterr = count*show_every\n",
    "                    # Show example output for the generator # 25 number for 1 time\n",
    "                    images_grid = show_generator_output(sess, 25, input_z, data_shape[3], data_image_mode)\n",
    "                    x = 0\n",
    "                    for image_grid in images_grid : \n",
    "                        x = x+1\n",
    "                        dst = os.path.join(\"output\", str(steps),str(iterr)+str(x)+\".png\")\n",
    "                        pyplot.imsave(dst, image_grid)\n",
    "                        \n",
    "                 # saving the model         \n",
    "                if epoch_i % 10 == 0:\n",
    "                    if not os.path.exists('./model/'):\n",
    "                        os.makedirs('./model')\n",
    "                    saver.save(sess, './model/' + str(epoch_i))      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "140\n",
      "(?, 4, 4, 1024)\n",
      "(?, 6, 6, 512)\n",
      "(?, 12, 12, 256)\n",
      "(?, 25, 25, 3)\n",
      "INFO:tensorflow:Restoring parameters from /tmp/model.ckpt\n"
     ]
    },
    {
     "ename": "FileExistsError",
     "evalue": "[Errno 17] File exists: 'output'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileExistsError\u001b[0m                           Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-3cf64f8b526a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0mceleba_dataset\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDataset\u001b[0m\u001b[0;34m(\u001b[0m \u001b[0mglob\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'./motionpatch/*.png'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGraph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_default\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepochs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mz_dim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlearning_rate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbeta1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mceleba_dataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_batches\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mceleba_dataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mceleba_dataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimage_mode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-8-4eafe8fdaf6d>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(epoch_count, batch_size, z_dim, learning_rate, beta1, get_batches, data_shape, data_image_mode, print_every, show_every)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0mthreads\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart_queue_runners\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msess\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msess\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcoord\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcoord\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m         \u001b[0;31m#sess.run(tf.global_variables_initializer())\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m         \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmkdir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'output'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mepoch_i\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepoch_count\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mbatch_images\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mget_batches\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileExistsError\u001b[0m: [Errno 17] File exists: 'output'"
     ]
    }
   ],
   "source": [
    "batch_size = 50\n",
    "z_dim = 100\n",
    "learning_rate = 0.00025\n",
    "beta1 = 0.45\n",
    "\n",
    "epochs = 500\n",
    "print(len(glob('./motionpatch/*.png')))\n",
    "celeba_dataset = Dataset( glob('./motionpatch/*.png'))\n",
    "with tf.Graph().as_default():\n",
    "    train(epochs, batch_size, z_dim, learning_rate, beta1, celeba_dataset.get_batches, celeba_dataset.shape, celeba_dataset.image_mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}