Super-Mega-Model-128x128.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Use GPU No 0.\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))\n",
    "\n",
    "from keras.layers.convolutional import UpSampling2D, Conv2D, MaxPooling2D\n",
    "from keras.optimizers  import Adam, Adagrad\n",
    "from keras.activations import sigmoid,relu\n",
    "from keras.callbacks   import TensorBoard\n",
    "from keras.models import Model, Sequential\n",
    "from keras.utils  import to_categorical\n",
    "from keras.models import load_model\n",
    "from keras.layers import Input, concatenate, multiply, GaussianNoise, Flatten, Dense, Lambda, Subtract, Activation, multiply, Add, Conv2DTranspose, AveragePooling2D, MaxPooling2D\n",
    "from keras import backend as K\n",
    "import keras\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import time\n",
    "from Datasets import Dataset, get_model_inputs\n",
    "from perceptual_model import get_perceptual_model\n",
    "\n",
    "import os\n",
    "import tensorflow as tf\n",
    "from PIL import ImageFile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#size of input image: c*c\n",
    "c = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_models(savekey, epoch_reg, save_disc=True, save_gen=True):\n",
    "    if save_disc:\n",
    "        dis_model.save(\"../../data/team6/model_checkpoints/{}_dis_model{}.h5\".format(savekey,epoch_reg))\n",
    "    if save_gen:\n",
    "        generator.save(\"../../data/team6/model_checkpoints/{}_gen_model{}.h5\".format(savekey,epoch_reg))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def MRU(X,I, filter_depth, deconv=False):\n",
    "    \n",
    "    print(\"Input sizes: \",I.shape,X.shape)\n",
    "\n",
    "    out_size = X.get_shape().as_list()[-1]\n",
    "    new_size = X.get_shape().as_list()[-1]\n",
    "    \n",
    "    #deconv or conv\n",
    "    if deconv:\n",
    "        #same as upsample\n",
    "        X = Conv2DTranspose(filters = out_size, kernel_size=(2,2), strides=(2,2))(X)\n",
    "\n",
    "    merge_one        = concatenate([X,I])\n",
    "\n",
    "    #mi\n",
    "    conv_sig_m     = Conv2D(filters = out_size, kernel_size=(3, 3), padding=\"same\", activation=\"sigmoid\")(merge_one)\n",
    "    \n",
    "    #ni\n",
    "    conv_sig_n       = Conv2D(filters = filter_depth, kernel_size=(3, 3), padding=\"same\", activation=\"sigmoid\")(merge_one)\n",
    "    \n",
    "    mul_lower        = multiply([conv_sig_m,X])\n",
    "\n",
    "    merge_two        = concatenate([mul_lower,I])\n",
    "\n",
    "    conv_func_one    = Conv2D(filters = filter_depth, kernel_size=(3, 3), padding=\"same\", activation=\"relu\")(merge_two)\n",
    "    \n",
    "    mul_higher       = multiply([conv_sig_n, conv_func_one])\n",
    "\n",
    "    conv_sig_sub     = Lambda(lambda x: 1 - x)(conv_sig_n)\n",
    "    \n",
    "    conv_x           = Conv2D(filters=filter_depth, kernel_size=(3, 3), padding=\"same\", activation=\"relu\")(X)\n",
    "\n",
    "    mul_final        = multiply([conv_x,conv_sig_sub])\n",
    "\n",
    "    sum_final        = Add()([mul_higher,mul_final])\n",
    "    \n",
    "    if not deconv:\n",
    "        sum_final    = MaxPooling2D(pool_size = 2)(sum_final)\n",
    "       \n",
    "    #my addition\n",
    "    sum_final        = Activation(\"relu\")(sum_final)  \n",
    "    \n",
    "    return sum_final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_xx_model(img_shape=(c,c,3)):\n",
    "    I = Input(shape=img_shape)\n",
    "    \n",
    "    out = Conv2D(filters=9, kernel_size=(c,c))(I)\n",
    "\n",
    "    return Model(I, out, name='xx_layer')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_encoder(I):\n",
    "\n",
    "    K1  = concatenate([I,I])\n",
    "    new = MRU(I  ,I ,32)\n",
    "    I2  = AveragePooling2D(pool_size=2)(I)\n",
    "    \n",
    "    K2  = concatenate([new,I2])\n",
    "    new = MRU(new,I2,64)\n",
    "    I3  = AveragePooling2D(pool_size=2)(I2)\n",
    "    \n",
    "    K3  = concatenate([new,I3])\n",
    "    new = MRU(new,I3,128)\n",
    "    I4  = AveragePooling2D(pool_size=2)(I3)\n",
    "\n",
    "    K4  = concatenate([new,I4])\n",
    "    out = MRU(new,I4,256)\n",
    "    I5  = AveragePooling2D(pool_size=2)(I4)\n",
    "    \n",
    "    K5  = concatenate([out,I5])\n",
    "    \n",
    "    encoder = Model(I, out, name='encoder');\n",
    "    \n",
    "    return encoder,K4,K3,K2,K1\n",
    "    \n",
    "def get_decoder(feature_map, K4,K3,K2,K1,I):\n",
    "\n",
    "    decoder_in = Input(shape=(int(feature_map.shape[1]),int(feature_map.shape[2]),int(feature_map.shape[3])))\n",
    "\n",
    "    new = MRU(decoder_in,K4,128,deconv=True)\n",
    "    new = MRU(new,K3,64,deconv=True)\n",
    "    new = MRU(new,K2,32,deconv=True)\n",
    "    out = MRU(new,K1, 3,deconv=True)\n",
    "    \n",
    "    decoder = Model([decoder_in,I], out, name='decoder');\n",
    "    \n",
    "    return decoder\n",
    "\n",
    "def get_generator(img_shape, color_info_shape):\n",
    "    \n",
    "    I = Input(shape=img_shape)\n",
    "    R = Input(shape=color_info_shape)\n",
    "    \n",
    "    encoder,K4,K3,K2,K1    = get_encoder(I)\n",
    "    \n",
    "    decoder    = get_decoder(encoder.output, K4,K3,K2,K1,I)\n",
    "    \n",
    "    encoded    = encoder(I)\n",
    "    \n",
    "    xy_layer   = Conv2DTranspose(filters=9, kernel_size=(8,8))\n",
    "    \n",
    "    colors     = xy_layer(R)\n",
    "    \n",
    "    encolor    = concatenate([colors,encoded])\n",
    "\n",
    "    out        = decoder([encoded,I])\n",
    "    \n",
    "    return Model([I,R], out, name='generator_model')\n",
    "\n",
    "def get_discriminator(img_shape):\n",
    "    \n",
    "    I = Input(shape=img_shape)\n",
    "    out = MRU(I, I, 8)\n",
    "\n",
    "    return Model(I, out, name='discriminator_skeleton')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create datasets\n",
    "\n",
    "dataset_parent_path = os.path.join('..','..','data','team6','celeba_ydk')\n",
    "dataset = Dataset(\"celeba\", dataset_parent_path, False, 1)\n",
    "save_key = \"celeba-ct\"\n",
    "\n",
    "num_class = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GAN parameters\n",
    "img_rows = c\n",
    "img_cols = c\n",
    "channels = 3\n",
    "img_shape = (img_rows, img_cols, channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_objects = dataset.get_all_image_objects()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "percept_model = get_perceptual_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#constructing xx_model\n",
    "\n",
    "in_real      = Input(shape=(img_shape))\n",
    "\n",
    "xx_layer     = Conv2D(filters=9, kernel_size=(128,128), name=\"xx_model\", kernel_initializer=\"random_uniform\")\n",
    "\n",
    "color_inf1   = xx_layer(in_real)\n",
    "\n",
    "xx_model     = Model(input=in_real, outputs=color_inf1)\n",
    "\n",
    "# Compiling generator model\n",
    "losses = {'xx_model': 'mse'}\n",
    "\n",
    "lossWeights = {\"xx_model\": 1.0}\n",
    "\n",
    "\n",
    "opt = keras.optimizers.Adam(lr=2e-4)\n",
    "xx_model.compile(optimizer=opt, loss=losses, loss_weights=lossWeights, metrics=['mse'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constructing generator\n",
    "\n",
    "\n",
    "# Generator Input\n",
    "in_gen = Input(shape=(img_shape))\n",
    "#color map (random noise at prediction)\n",
    "in_Rmp = Input(shape=(1,1,9))\n",
    "\n",
    "\n",
    "# Get Generator model\n",
    "generator = get_generator(img_shape,(1,1,9))\n",
    "\n",
    "# Normal output of generator\n",
    "out_gen = generator([in_gen,in_Rmp])\n",
    "\n",
    "# Final generator model\n",
    "gen_model = Model([in_gen,in_Rmp], out_gen)\n",
    "\n",
    "# Constructing discriminator\n",
    "\n",
    "# Discriminator input \n",
    "disc_in = Input(shape=img_shape)\n",
    "\n",
    "# Get skeleton of discriminator\n",
    "disc = get_discriminator(img_shape)\n",
    "\n",
    "# Pass generator output \n",
    "x = disc(disc_in)\n",
    "\n",
    "# Create output layers\n",
    "validity_dense = Dense(1, activation='sigmoid', name='end_sigmoid')\n",
    "classification_dense = Dense(num_class, activation='softmax', name='end_softmax')\n",
    "\n",
    "flat_white = Flatten()(x)\n",
    "\n",
    "dis_valid_dense = validity_dense(flat_white)\n",
    "dis_class_dense = classification_dense(flat_white)\n",
    "\n",
    "# Final discriminator model\n",
    "dis_model = Model(disc_in, [\n",
    "                            dis_valid_dense, \n",
    "                            dis_class_dense\n",
    "                            ] )\n",
    "\n",
    "# Constructing GAN \n",
    "\n",
    "# Create input layers\n",
    "in_gan = Input(shape=(img_shape))\n",
    "out_gen_for_gan = generator([in_gan,in_Rmp])\n",
    "\n",
    "# Pass generator output from perceptual model for its loss\n",
    "percept_out_for_gan = percept_model(out_gen_for_gan)\n",
    "\n",
    "# Pass generator output from discriminator model for\n",
    "gan_dis_sket_out = disc(out_gen_for_gan)\n",
    "\n",
    "# \n",
    "flat_black = Flatten()(gan_dis_sket_out)\n",
    "\n",
    "# Pass from validity dense layer\n",
    "gan_valid_out = validity_dense(flat_black)\n",
    "# Pass from auxilary classification dense layer\n",
    "gan_class_out = classification_dense(flat_black)\n",
    "\n",
    "# Final GAN model\n",
    "gan_model = Model(\n",
    "                  input=[in_gan, in_Rmp], \n",
    "                  outputs=[\n",
    "                      out_gen_for_gan, \n",
    "                           gan_valid_out, \n",
    "                           gan_class_out, \n",
    "                           percept_out_for_gan,\n",
    "                          ]\n",
    "                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compiling gan model\n",
    "losses = {\n",
    "    'perceptual-model': 'mse',\n",
    "    'generator_model': 'mse',\n",
    "    'end_softmax': 'categorical_crossentropy',\n",
    "    'end_sigmoid': 'binary_crossentropy',\n",
    "    }\n",
    "\n",
    "lossWeights = {\"end_softmax\": -0.5, # Want to maximize it in GAN\n",
    "               \"end_sigmoid\": 0.5,\n",
    "               \"generator_model\": 0.5,\n",
    "               \"perceptual-model\": 0.01,\n",
    "              }\n",
    "\n",
    "# Make discriminator not trainable for gan\n",
    "\n",
    "gan_model.layers[-3].trainable = False\n",
    "gan_model.layers[-2].trainable = False\n",
    "gan_model.layers[-1].trainable = False\n",
    "\n",
    "gan_opt = keras.optimizers.Adam(lr=2e-4)\n",
    "\n",
    "gan_model.compile(loss=losses, optimizer=gan_opt,loss_weights=lossWeights, metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compiling generator model\n",
    "losses = {'generator_model': 'mse'}\n",
    "\n",
    "lossWeights = {\"generator_model\": 0.1,}\n",
    "\n",
    "\n",
    "gen_opt = keras.optimizers.Adam(lr=2e-4)\n",
    "gen_model.compile(optimizer=gen_opt, loss=losses, loss_weights=lossWeights, metrics=['mse'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compiling discriminator model\n",
    "losses = {'end_sigmoid': 'mse',\n",
    "          'end_softmax': 'mse'\n",
    "         }\n",
    "\n",
    "lossWeights = {\"end_sigmoid\": 1.0, \n",
    "               \"end_softmax\": 1.0, \n",
    "              }\n",
    "\n",
    "dis_opt = keras.optimizers.Adam(lr=1e-6)\n",
    "\n",
    "dis_model.trainable = True\n",
    "dis_model.layers[1].layers[1].trainable = False\n",
    "dis_model.compile(loss=losses, optimizer=dis_opt,loss_weights=lossWeights, metrics=['mse'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save logs to ./v2_logs\n",
    "logdir = 'v2_logs'\n",
    "if not os.path.exists(logdir):\n",
    "    os.mkdir(logdir)\n",
    "callback = TensorBoard(logdir, write_graph=True)\n",
    "callback.set_model(gan_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_log(callback, names, logs, batch_no):\n",
    "    for name, value in zip(names, logs):\n",
    "        summary = tf.Summary()\n",
    "        summary_value = summary.value.add()\n",
    "        summary_value.simple_value = value\n",
    "        summary_value.tag = name\n",
    "        callback.writer.add_summary(summary, batch_no)\n",
    "        callback.writer.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save example predictions.\n",
    "\n",
    "def save_imgs(epoch, input_data, iteration,col=np.array([])):\n",
    "    r, c = 2, 2\n",
    "    idx = np.array([12,5,2,10]) \n",
    "    elements = input_data[idx,:,:]\n",
    "    if col.size<=1:\n",
    "        col = np.random.random_sample((elements.size,1,1,9))\n",
    "    gen_imgs = gen_model.predict([elements, col])\n",
    "\n",
    "    fig, axs = plt.subplots(r, c)\n",
    "    cnt = 0\n",
    "    for i in range(r):\n",
    "        for j in range(c):\n",
    "            axs[i,j].imshow(gen_imgs[cnt, :,:,:]/255)\n",
    "            axs[i,j].axis('off')\n",
    "            cnt += 1\n",
    "    fig.suptitle('EPOCH-ITERATION: {}-{}'.format(epoch, iteration))\n",
    "    fig.savefig(\"results/{}/{}_{}_{}.png\".format(save_key,save_key,epoch,iteration))\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resuming_training = False\n",
    "\n",
    "if resuming_training:\n",
    "    reg_all_epochs = epoch\n",
    "    goto_save = True\n",
    "    iter_reg = iter_num\n",
    "else:\n",
    "    reg_all_epochs = 0\n",
    "    goto_save = False\n",
    "    iter_reg = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "time_init_train = time.time()\n",
    "\n",
    "batch_size    = 32\n",
    "iter_count    = len(img_objects) // batch_size\n",
    "total_count   = 0\n",
    "initial_epoch = reg_all_epochs\n",
    "epochs        = 100\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    \n",
    "    reg_all_epochs = initial_epoch + epoch\n",
    "    \n",
    "    print(\"Enter epoch {}.\".format(reg_all_epochs));\n",
    "    \n",
    "    epoch_init_time = time.time()\n",
    "    \n",
    "    if goto_save:\n",
    "        iter_low  = iter_reg\n",
    "        goto_save = False\n",
    "    else:\n",
    "        iter_low  = 1\n",
    "\n",
    "    \n",
    "    for iter_num in range(iter_low, iter_count):\n",
    "        \n",
    "        total_count += 1\n",
    "        \n",
    "        # Obtain batch of inputs\n",
    "        batch_objects = img_objects[batch_size*iter_num:batch_size*(iter_num+1)]\n",
    "        #try:\n",
    "        in_x0, in_x1, class_labels = get_model_inputs(batch_objects, img_shape, num_class)\n",
    "        \n",
    "        if epoch ==0 and iter_num == 1:\n",
    "            #get a sample for example imgs\n",
    "            sv_imgs_in = in_x0\n",
    "\n",
    "        # At every 100 iteration, save example predictions.\n",
    "        if iter_num%100 ==0:\n",
    "            save_imgs(reg_all_epochs, sv_imgs_in, iter_num)\n",
    "            print(\"{}/{}\".format(iter_num,iter_count) )\n",
    "        colors = xx_model.predict(in_x1)\n",
    "        \n",
    "        # Forward pass the generator to get predicted examples\n",
    "        gen_sketches = gen_model.predict([in_x0,colors])\n",
    "        \n",
    "        # Pass real inputs from perceptual model to be use in loss\n",
    "        org_percepts = percept_model.predict(in_x1)\n",
    "    \n",
    "\n",
    "        # train the discriminator with class 1 only\n",
    "        if epoch + iter_num %2 == 0:\n",
    "            disc_losses = dis_model.train_on_batch(gen_sketches, [np.zeros(batch_size), class_labels])\n",
    "\n",
    "        # train the discriminator with class 2 only\n",
    "        else:\n",
    "            disc_losses = dis_model.train_on_batch(in_x1, [np.ones(batch_size), class_labels] )\n",
    "\n",
    "        gan_losses = gan_model.train_on_batch([in_x0,colors], [in_x1, np.ones(batch_size), \n",
    "                                                      class_labels,\n",
    "                                                      org_percepts,\n",
    "                                                     ])\n",
    "                                                          \n",
    "        \n",
    "        write_log(callback, gan_model.metrics_names, gan_losses, total_count)\n",
    "        \n",
    "        #save model every 2000 iterations.\n",
    "        if(iter_num%2000==0):\n",
    "            save_models(save_key, reg_all_epochs, save_disc=True, save_gen=True)\n",
    "            print(\"\\tTotal time elapsed: {:.2f} secs.\".format(time.time()-time_init_train))\n",
    "        \n",
    "    \n",
    "    print(\"Epoch {} finished in {:.2f} secs.\".format(reg_all_epochs, time.time()-epoch_init_time))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    " \n",
    "# Predict from source files. \n",
    "# gt_bit 1 if ground-truth exists, 0 otherwise.\n",
    "\n",
    "def predict_real(addr, interpolation = 0, gt_bit=0):\n",
    "\n",
    "    img_original  = cv2.imread(addr)\n",
    "\n",
    "    img  = cv2.resize(img_original, (c,c), interpolation=interpolation) \n",
    "    imgx = np.expand_dims(img,0)\n",
    "    img_result = gen_model.predict([imgx,np.random.random_sample((1,1,1,9))])\n",
    "    img_result = np.array(img_result[0])\n",
    "    img_result = np.reshape(img_result, img_result.shape[:])/255\n",
    "    fig, axs = plt.subplots(1, 2+gt_bit)\n",
    "    \n",
    "    axs[0].imshow(img)\n",
    "    axs[1].imshow(img_result)\n",
    "    axs[1].axis('off')\n",
    "    if gt_bit==1:\n",
    "        tmp = addr.split(\"/\")\n",
    "        tmp[-2] = \"data\"\n",
    "        tmp = str.join(\"/\",tmp)\n",
    "        img_gt  = cv2.imread(tmp)\n",
    "        img_gt  = cv2.resize(img_gt, (c,c))\n",
    "        axs[2].imshow(img_gt[...,::-1])\n",
    "        axs[2].axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save source images\n",
    "\n",
    "def save_src(input_data=sv_imgs_in):\n",
    "    r, c = 2, 2\n",
    "    idx = np.array([12,4,2,11]) \n",
    "    elements = input_data[idx,:,:]\n",
    "    gen_imgs,_ = gen_model.predict(elements)\n",
    "\n",
    "    fig, axs = plt.subplots(r, c)\n",
    "    cnt = 0\n",
    "    for i in range(r):\n",
    "        for j in range(c):\n",
    "            axs[i,j].imshow(input_data[idx[i*r+j]]/255)\n",
    "            axs[i,j].axis('on')\n",
    "            cnt += 1\n",
    "    fig.suptitle('SOURCE EDGES')\n",
    "    fig.show()\n",
    "    plt.show()\n",
    "    fig.savefig(\"results/{}/src.png\".format(save_key))\n",
    "    plt.close()"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}