UNIT_tf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unsupervised Image-image Translation(UNIT) TensorFlow Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from collections import OrderedDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Funciton to chose the activation function\n",
    "def activate(linear, activation='leaky_relu'):\n",
    "        if activation == 'sigmoid':\n",
    "            return tf.nn.sigmoid(linear)\n",
    "        elif activation == 'softmax':\n",
    "            return tf.nn.softmax(linear)\n",
    "        elif activation == 'tanh':\n",
    "            return tf.nn.tanh(linear)\n",
    "        elif activation == 'relu':\n",
    "            return tf.nn.relu(linear)\n",
    "        elif activation == 'leaky_relu':\n",
    "            return tf.nn.leaky_relu(linear)\n",
    "        elif activation == 'linear':\n",
    "            return linear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_noise_layer(input_layer, std=1.0):\n",
    "    with tf.name_scope(\"gaussian_layer\"):\n",
    "        noise = tf.random_normal(shape=tf.shape(input_layer), mean=0.0, stddev=std, dtype=tf.float32) \n",
    "        return tf.add(input_layer, noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_normalization_layer(input_layer):\n",
    "    with tf.name_scope(\"batch_norm_layer\"):\n",
    "        dimension = input_layer.get_shape().as_list()[-1]\n",
    "\n",
    "        mean, variance = tf.nn.moments(input_layer, axes=[0, 1, 2])\n",
    "\n",
    "        beta = tf.Variable(tf.constant(0.0, shape=[dimension]))\n",
    "\n",
    "        gamma = tf.Variable(tf.constant(1.0, shape=[dimension]))\n",
    "\n",
    "        bn_layer = tf.nn.batch_normalization(input_layer, mean, variance, beta, gamma, variance_epsilon=1e-10)\n",
    "\n",
    "        return bn_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_conv_layer(input_layer,     # The previous layer.\n",
    "                   filter_size,        # Width and height of each filter.\n",
    "                   num_filters,        # Number of filters.\n",
    "                   use_pooling=False,  \n",
    "                   pad = [],\n",
    "                   strides=[1,1,1,1],\n",
    "                   deconv=False,\n",
    "                   out_shape = [],      # Output shape in case of deconv \n",
    "                   batch_normalization=False,\n",
    "                   activation='leaky_relu'): # Use 2x2 max-pooling.\n",
    "\n",
    "    \n",
    "    num_input_channels = input_layer.get_shape().as_list()[-1]\n",
    "    # Shape of the filter-weights for the convolution. \n",
    "    # This format is determined by the TensorFlow API.\n",
    "    if deconv:\n",
    "        shape = [filter_size, filter_size, num_filters, num_input_channels]\n",
    "    else:\n",
    "        shape = [filter_size, filter_size, num_input_channels, num_filters]\n",
    "\n",
    "    # Create new weights aka. filters with the given shape.\n",
    "    weights = tf.Variable(tf.truncated_normal(shape, stddev=0.05))\n",
    "\n",
    "    # Create new biases, one for each filter.\n",
    "    biases = tf.Variable(tf.constant(1.0, shape=[num_filters]))\n",
    "\n",
    "    if len(pad) > 0:\n",
    "        input_layer = tf.pad(input_layer, [[0, 0], [pad[0], pad[0]], [pad[1], pad[1]], [0, 0]], \"CONSTANT\")\n",
    "    else:\n",
    "        pad = [0,0]\n",
    "    # Create the TensorFlow operation for de-convolution.\n",
    "    if deconv:\n",
    "        #in_shape = input_layer.get_shape().as_list()\n",
    "        in_shape = tf.shape(input_layer)\n",
    "    \n",
    "        out_h = ((in_shape[1] - 1) * strides[1]) + filter_size - 2 * pad[0]\n",
    "        \n",
    "        out_w = ((in_shape[2] - 1) * strides[2]) + filter_size - 2 * pad[1]\n",
    "        \n",
    "        output_shape = tf.stack([in_shape[0], out_h, out_w, num_filters])\n",
    "        \n",
    "        layer = tf.nn.conv2d_transpose(value=input_layer, \n",
    "                                       filter=weights, \n",
    "                                       output_shape=output_shape, \n",
    "                                       strides=strides, \n",
    "                                       padding='SAME')\n",
    "        \n",
    "    else:\n",
    "        layer = tf.nn.conv2d(input=input_layer,\n",
    "                         filter=weights,\n",
    "                         strides=strides,\n",
    "                         padding='SAME')\n",
    "\n",
    "    # Add the biases to the results of the convolution.\n",
    "    # A bias-value is added to each filter-channel.\n",
    "    layer = tf.add(layer,biases)\n",
    "\n",
    "    # Use pooling to down-sample the image resolution?\n",
    "    if use_pooling:\n",
    "        # This is 2x2 max-pooling, which means that we\n",
    "        # consider 2x2 windows and select the largest value\n",
    "        # in each window. Then we move 2 pixels to the next window.\n",
    "        layer = tf.nn.max_pool(value=layer,\n",
    "                               ksize=[1, 2, 2, 1],\n",
    "                               strides=[1, 2, 2, 1],\n",
    "                               padding='SAME')\n",
    "\n",
    "    # Batch normalization\n",
    "    if batch_normalization:\n",
    "        layer = batch_normalization_layer(layer)\n",
    "    \n",
    "    # Activation of the layers (ReLU).\n",
    "    layer = activate(layer, activation=activation)\n",
    "\n",
    "    return layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_res_block(X, filter_size=3, num_filters=3):\n",
    "    \n",
    "    num_filters = X.get_shape().as_list()[-1]\n",
    "    \n",
    "    layer_conv1 = create_conv_layer(input_layer=X,\n",
    "                                   filter_size=filter_size,\n",
    "                                   num_filters=num_filters,\n",
    "                                   use_pooling=False,\n",
    "                                   batch_normalization=True,\n",
    "                                   activation='relu')\n",
    "    \n",
    "    layer_conv2 = create_conv_layer(input_layer=layer_conv1,\n",
    "                                   filter_size=filter_size,\n",
    "                                   num_filters=num_filters,\n",
    "                                   use_pooling=False,\n",
    "                                   batch_normalization=False,\n",
    "                                   activation='linear')\n",
    "    \n",
    "    layer_conv2 = batch_normalization_layer(layer_conv2)\n",
    "    \n",
    "    layer_conv2 += X\n",
    "    \n",
    "    layer_res = activate(layer_conv2, activation='relu')\n",
    "    \n",
    "    return layer_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_encoder(X, layer_n, res_block_n, ly_name):\n",
    "    with tf.name_scope(ly_name):\n",
    "        num_filters = X.get_shape().as_list()[-1]\n",
    "\n",
    "        encoder = OrderedDict()\n",
    "\n",
    "        encoder['layer_conv1'] = create_conv_layer(input_layer=X,\n",
    "                                           num_filters=num_filters,\n",
    "                                           filter_size=7,\n",
    "                                           strides=[1,1,1,1])\n",
    "\n",
    "        for i in range(1, layer_n):\n",
    "\n",
    "            encoder['layer_conv'+str(i+1)] = create_conv_layer(encoder[next(reversed(encoder))],\n",
    "                                           num_filters=num_filters*2,\n",
    "                                           filter_size=3,\n",
    "                                           strides=[1,2,2,1])\n",
    "\n",
    "            num_filters = num_filters*2\n",
    "\n",
    "        for i in range(0, res_block_n):\n",
    "            encoder['block_en_res'+str(i+1)] = create_res_block(encoder[next(reversed(encoder))],\n",
    "                                                               num_filters=num_filters)\n",
    "        \n",
    "    return encoder#[next(reversed(encoder))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_shared_layers(X, block_shared_n, ly_name):\n",
    "    with tf.name_scope(ly_name):\n",
    "        encoder_shared = OrderedDict()\n",
    "        decoder_shared = OrderedDict()\n",
    "\n",
    "        encoder_shared_ = create_res_block(X)\n",
    "        encoder_shared['block_shared_res1'] = gaussian_noise_layer(encoder_shared_)\n",
    "\n",
    "        for i in range(1, block_shared_n+1):\n",
    "            encoder_shared_ = create_res_block(encoder_shared[next(reversed(encoder_shared))])\n",
    "            encoder_shared['block_shared_res'+str(i+1)] = gaussian_noise_layer(encoder_shared_)\n",
    "\n",
    "        decoder_shared['block_shared_res1'] = create_res_block(encoder_shared[next(reversed(encoder_shared))])\n",
    "\n",
    "        for i in range(1, block_shared_n+1):\n",
    "            decoder_shared['block_shared_res'+str(i+1)] = create_res_block(decoder_shared[next(reversed(decoder_shared))])\n",
    "\n",
    "        return decoder_shared#[next(reversed(decoder_shared))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_generator(X, layer_n, res_block_n, out_shape, ly_name):\n",
    "    with tf.name_scope(ly_name):\n",
    "        num_filters = X.get_shape().as_list()[-1]\n",
    "\n",
    "        decoder = OrderedDict()\n",
    "\n",
    "        decoder['block_res1'] = create_res_block(X, num_filters=num_filters)\n",
    "\n",
    "        for i in range(1, res_block_n):\n",
    "            decoder['block_res'+str(i+1)] = create_res_block(decoder[next(reversed(decoder))],\n",
    "                                                               num_filters=num_filters)\n",
    "\n",
    "        for i in range(0, layer_n-1):\n",
    "\n",
    "            num_filters = num_filters//2\n",
    "\n",
    "            decoder['layer_deconv'+str(i+1)] = create_conv_layer(decoder[next(reversed(decoder))],\n",
    "                                           num_filters=num_filters,\n",
    "                                           filter_size=2,\n",
    "                                           strides=[1,2,2,1],\n",
    "                                           deconv=True,\n",
    "                                           batch_normalization=True,\n",
    "                                           activation='relu')\n",
    "\n",
    "        decoder['layer_deconv_final'] = create_conv_layer(decoder[next(reversed(decoder))],\n",
    "                                           num_filters=num_filters,\n",
    "                                           filter_size=1,\n",
    "                                           strides=[1,1,1,1],\n",
    "                                           deconv=True,\n",
    "                                           activation='tanh')\n",
    "\n",
    "        return decoder#[next(reversed(decoder))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_discrimiator(X, num_filters, layer_n, ly_name):\n",
    "    with tf.name_scope(ly_name):\n",
    "        discriminator = OrderedDict()\n",
    "\n",
    "        discriminator['layer_discrim_conv1'] = create_conv_layer(input_layer=X,\n",
    "                                           num_filters=num_filters,\n",
    "                                           filter_size=3,\n",
    "                                           strides=[1,2,2,1])\n",
    "\n",
    "        for i in range(1, layer_n):\n",
    "\n",
    "            num_filters = num_filters * 2\n",
    "\n",
    "            discriminator['layer_discrim_convb'+str(i+1)] = create_conv_layer(discriminator[next(reversed(discriminator))],\n",
    "                                           num_filters=num_filters,\n",
    "                                           filter_size=3,\n",
    "                                           strides=[1,2,2,1])\n",
    "\n",
    "        return discriminator#[next(reversed(discriminator))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reshape_img(x,o_h=484,o_w=642):\n",
    "    return tf.image.resize_image_with_crop_or_pad(image=x,\n",
    "                                                  target_height=o_h,\n",
    "                                                  target_width=o_w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_encoding_loss(x):\n",
    "    x_2 = tf.pow(x, 2)\n",
    "    encoding_loss = tf.reduce_mean(x_2)\n",
    "    return encoding_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_l1_loss(x,y):\n",
    "    return tf.reduce_mean(tf.subtract(x,y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _compute_true_acc(predictions):\n",
    "    predictions = tf.cast(tf.greater_equal(predictions, 0.5), dtype=tf.float32)\n",
    "    acc = tf.reduce_sum(predictions) / tf.cast(tf.shape(predictions)[1], dtype=tf.float32)\n",
    "    return acc\n",
    "\n",
    "def _compute_fake_acc(predictions):\n",
    "    predictions = tf.cast(tf.less(predictions, 0.5), dtype=tf.float32)\n",
    "    acc = tf.reduce_sum(predictions) / tf.cast(tf.shape(predictions)[1], dtype=tf.float32)\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Build Graph\n",
    "tf.reset_default_graph()\n",
    "\n",
    "X_A = tf.placeholder(tf.float32, shape=[1, None, None, 3], name='X_A')\n",
    "X_B = tf.placeholder(tf.float32, shape=[1, None, None, 3], name='X_B')\n",
    "\n",
    "encode_A = create_encoder(X_A, layer_n=3, res_block_n=3, ly_name='encoder_A')\n",
    "encode_B = create_encoder(X_B, layer_n=3, res_block_n=3, ly_name='encoder_B')\n",
    "\n",
    "encode_AB = tf.concat([encode_A[next(reversed(encode_A))], encode_B[next(reversed(encode_B))]], axis=1)\n",
    "\n",
    "shared_in = tf.placeholder(tf.float32, shape=encode_AB.get_shape().as_list(), name='shared_in')\n",
    "shared_ly = create_shared_layers(shared_in, block_shared_n=3, ly_name='shared_ly')\n",
    "\n",
    "decode_ch = shared_ly[next(reversed(shared_ly))].get_shape().as_list()[-1]\n",
    "\n",
    "decode_A_in = tf.placeholder(tf.float32, shape=[1, None, None, decode_ch], name='decode_A_in')\n",
    "decode_A = create_generator(decode_A_in, layer_n=3, res_block_n=3, out_shape=[1, None, None, 3], ly_name='decoder_A')\n",
    "\n",
    "decode_B_in = tf.placeholder(tf.float32, shape=[1, None, None, decode_ch], name='decode_B_in')\n",
    "decode_B = create_generator(decode_B_in, layer_n=3, res_block_n=3, out_shape=[1, None, None, 3], ly_name='decoder_B')\n",
    "\n",
    "X_aa, X_ba = tf.split(decode_A[next(reversed(decode_A))], \n",
    "                      num_or_size_splits=2, axis=1)\n",
    "X_ab, X_bb = tf.split(decode_B[next(reversed(decode_B))], \n",
    "                      num_or_size_splits=2, axis=1)\n",
    "\n",
    "X_aa, X_ba, X_ab, X_bb = [reshape_img(x) for x in [X_aa, X_ba, X_ab, X_bb]]\n",
    "\n",
    "#dis\n",
    "\n",
    "X_ba_in = tf.placeholder(tf.float32, shape=[1, None, None, 3], name='X_ba_in')\n",
    "X_ab_in = tf.placeholder(tf.float32, shape=[1, None, None, 3], name='X_ab_in')\n",
    "\n",
    "outs_a = create_discrimiator(X_ba_in, num_filters=3, layer_n=6, ly_name='discriminator_A')\n",
    "outs_b = create_discrimiator(X_ab_in, num_filters=3, layer_n=6, ly_name='discriminator_B')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator_update(sess, images_a, images_b, hyperparameters):\n",
    "    \n",
    "    encode_ab = sess.run(encode_AB, feed_dict={X_A:images_a, \n",
    "                                            X_B:images_b})\n",
    "    \n",
    "    shared = sess.run(shared_ly[next(reversed(shared_ly))], feed_dict={shared_in:encode_ab})\n",
    "    \n",
    "    x_aa, x_ba, x_ab, x_bb = sess.run([X_aa, X_ba, X_ab, X_bb], \n",
    "                                               feed_dict={decode_A_in:shared,\n",
    "                                                          decode_B_in:shared})\n",
    "    \n",
    "    data_a = sess.run(tf.concat([images_a, x_ba], axis=1))\n",
    "    data_b = sess.run(tf.concat([images_b, x_ab], axis=1))\n",
    "    \n",
    "    \n",
    "    out_a, out_b = sess.run([outs_a[next(reversed(outs_a))], outs_b[next(reversed(outs_b))]],\n",
    "                            feed_dict={X_ba_in:data_a, X_ab_in:data_b})\n",
    "    \n",
    "    out_a = tf.contrib.layers.flatten(out_a)\n",
    "    out_b = tf.contrib.layers.flatten(out_b)\n",
    "    \n",
    "    with tf.name_scope('discriminator_cost'):\n",
    "\n",
    "        out_true_a, out_fake_a = tf.split(out_a, num_or_size_splits=2, axis=1)\n",
    "\n",
    "        out_true_b, out_fake_b = tf.split(out_b, num_or_size_splits=2, axis=1)\n",
    "\n",
    "        all1 = tf.Variable(tf.constant(1.0, shape=out_true_a.get_shape()))\n",
    "        all0 = tf.Variable(tf.constant(0.0, shape=out_fake_a.get_shape()))\n",
    "\n",
    "        ad_true_loss_a = tf.nn.sigmoid_cross_entropy_with_logits(logits=out_true_a, labels=all1)\n",
    "        ad_true_loss_b = tf.nn.sigmoid_cross_entropy_with_logits(logits=out_true_b, labels=all1)\n",
    "        ad_fake_loss_a = tf.nn.sigmoid_cross_entropy_with_logits(logits=out_fake_a, labels=all0)\n",
    "        ad_fake_loss_b = tf.nn.sigmoid_cross_entropy_with_logits(logits=out_fake_b, labels=all0)\n",
    "\n",
    "        ad_loss_a = tf.reduce_sum(ad_true_loss_a + ad_fake_loss_a)\n",
    "        ad_loss_b = tf.reduce_sum(ad_true_loss_b + ad_fake_loss_b)\n",
    "\n",
    "        true_a_acc,true_b_acc,fake_a_acc,fake_b_acc = sess.run([_compute_true_acc(out_true_a),\n",
    "                                                                _compute_true_acc(out_true_b),\n",
    "                                                                _compute_fake_acc(out_fake_a),\n",
    "                                                                _compute_fake_acc(out_fake_b)])\n",
    "\n",
    "        print(\"True Acc:\",0.5*(true_a_acc+true_b_acc),\" Fake Acc:\",0.5*(fake_a_acc+fake_b_acc))\n",
    "\n",
    "        cost = hyperparameters['gan_w'] * ( ad_loss_a + ad_loss_b )\n",
    "    \n",
    "        optimizer = tf.train.AdamOptimizer(learning_rate=0.0001,beta1=0.5,beta2=0.999).minimize(cost)\n",
    "    \n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    \n",
    "    _, loss = sess.run([optimizer,cost])\n",
    "    \n",
    "    print('Discriminator Cost: ',loss)\n",
    "                       \n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator_update(sess, images_a, images_b, hyperparameters):\n",
    "    \n",
    "    encode_ab = sess.run(encode_AB, feed_dict={X_A:images_a, \n",
    "                                            X_B:images_b})\n",
    "    \n",
    "    shared = sess.run(shared_ly[next(reversed(shared_ly))], feed_dict={shared_in:encode_ab})\n",
    "    \n",
    "    x_aa, x_ba, x_ab, x_bb = sess.run([X_aa, X_ba, X_ab, X_bb], \n",
    "                                               feed_dict={decode_A_in:shared,\n",
    "                                                          decode_B_in:shared})\n",
    "    #a2b\n",
    "    x_bab = sess.run(encode_A[next(reversed(encode_A))], feed_dict={X_A:x_ba})\n",
    "    shared_bab = sess.run(shared_ly[next(reversed(shared_ly))], feed_dict={shared_in:x_bab})\n",
    "    x_bab = sess.run(reshape_img(decode_B[next(reversed(decode_B))]), feed_dict={decode_B_in:shared_bab})\n",
    "    \n",
    "    #b2a\n",
    "    x_aba = sess.run(encode_B[next(reversed(encode_B))], feed_dict={X_B:x_ab})\n",
    "    shared_aba = sess.run(shared_ly[next(reversed(shared_ly))], feed_dict={shared_in:x_aba})\n",
    "    x_aba = sess.run(reshape_img(decode_A[next(reversed(decode_A))]), feed_dict={decode_A_in:shared_aba})\n",
    "    \n",
    "    #dis\n",
    "    out_a, out_b = sess.run([outs_a[next(reversed(outs_a))], outs_b[next(reversed(outs_b))]],\n",
    "                            feed_dict={X_ba_in:x_ba, X_ab_in:x_ab})\n",
    "    \n",
    "    out_a = tf.contrib.layers.flatten(out_a)\n",
    "    out_b = tf.contrib.layers.flatten(out_b)\n",
    "    \n",
    "    #loss\n",
    "    all_ones = tf.Variable(tf.constant(1.0, shape=out_a.get_shape()))\n",
    "    \n",
    "    with tf.name_scope('generator_cost'):\n",
    "        ad_loss_a = tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(logits=out_a, labels=all_ones))\n",
    "        ad_loss_b = tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(logits=out_b, labels=all_ones))\n",
    "\n",
    "        enc_loss  = compute_encoding_loss(shared)\n",
    "        enc_bab_loss = compute_encoding_loss(shared_bab)\n",
    "        enc_aba_loss = compute_encoding_loss(shared_aba)\n",
    "        ll_loss_a = compute_l1_loss(x_aa, images_a)\n",
    "        ll_loss_b = compute_l1_loss(x_bb, images_b)\n",
    "        ll_loss_aba = compute_l1_loss(x_aba, images_a)\n",
    "        ll_loss_bab = compute_l1_loss(x_bab, images_b)\n",
    "\n",
    "        total_loss = hyperparameters['gan_w'] * (ad_loss_a + ad_loss_b) + \\\n",
    "                     hyperparameters['ll_direct_link_w'] * (ll_loss_a + ll_loss_b) + \\\n",
    "                     hyperparameters['ll_cycle_link_w'] * (ll_loss_aba + ll_loss_bab) + \\\n",
    "                     hyperparameters['kl_direct_link_w'] * (enc_loss + enc_loss) + \\\n",
    "                     hyperparameters['kl_cycle_link_w'] * (enc_bab_loss + enc_aba_loss)\n",
    "\n",
    "        optimizer = tf.train.AdamOptimizer(learning_rate=0.0001, beta1=0.5,beta2=0.999).minimize(total_loss)\n",
    "    \n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    \n",
    "    _, loss = sess.run([optimizer,total_loss])\n",
    "    \n",
    "    print('Total_loss: ',loss)\n",
    "                       \n",
    "    return x_aa, x_ba, x_ab, x_bb, x_aba, x_bab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epoch = 100\n",
    "hyperparameters = {}\n",
    "\n",
    "hyperparameters['gan_w'] = 10\n",
    "hyperparameters['ll_direct_link_w'] = 100\n",
    "hyperparameters['ll_cycle_link_w'] = 100\n",
    "hyperparameters['kl_direct_link_w'] = 0.1\n",
    "hyperparameters['kl_cycle_link_w'] = 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import numpy as np\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_dir = \"/home/praveen/Downloads/ezgif-4-57d76e4278-gif-jpg\"\n",
    "\n",
    "images_names = os.listdir(images_dir)\n",
    "\n",
    "img_names = []    \n",
    "\n",
    "for file in images_names:\n",
    "    img_name = os.path.join(images_dir, file)\n",
    "    img_names.append(img_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  0\n",
      "True Acc: 0.703184187412  Fake Acc: 0.297940343618\n",
      "Discriminator Cost:  293001.0\n",
      "Total_loss:  29781.1\n",
      "Epoch:  1\n",
      "True Acc: 0.698508501053  Fake Acc: 0.288648188114\n",
      "Discriminator Cost:  304676.0\n",
      "Total_loss:  29619.1\n"
     ]
    }
   ],
   "source": [
    "sess = tf.InteractiveSession()\n",
    "    \n",
    "sess.run(tf.global_variables_initializer())\n",
    "\n",
    "for ep in range(2):\n",
    "    print('Epoch: ',ep)\n",
    "    for img_name in img_names[0:1]:\n",
    "        img = cv2.imread(img_name)\n",
    "        images_a, images_b = np.reshape(np.array(img[:,:642,:], dtype=\"float32\"),(1,484, 642, 3)),  \\\n",
    "                                np.reshape(np.array(img[:,642:,:], dtype=\"float32\"),(1,484, 642, 3))\n",
    "\n",
    "        discriminator_update(sess, images_a, images_b, hyperparameters)\n",
    "        x_aa, x_ba, x_ab, x_bb, x_aba, x_bab = generator_update(sess, images_a, images_b, hyperparameters)\n",
    "\n",
    "writer = tf.summary.FileWriter(\"./logs/\", sess.graph)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}