Fully Convolutional Variational Auto Encoders.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "80 Auto Seg d",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/hrmello/PhD/blob/master/Fully%20Convolutional%20Variational%20Auto%20Encoders.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DAHYWTA-HIa3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import sys, os,cv2\n",
        "from sklearn.utils import shuffle\n",
        "from scipy.misc import imread,imresize\n",
        "import matplotlib.pyplot as plt\n",
        "from sklearn.preprocessing import OneHotEncoder\n",
        "from skimage.transform import resize\n",
        "from imgaug import augmenters as iaa\n",
        "import imgaug as ia\n",
        "from skimage.color import rgba2rgb\n",
        "from skimage.color import rgb2gray\n",
        "import matplotlib\n",
        "\n",
        "plt.style.use('seaborn-white')\n",
        "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
        "np.random.seed(6278)\n",
        "tf.set_random_seed(6728)\n",
        "ia.seed(6278)\n",
        "\n",
        "def tf_elu(x): return tf.nn.elu(x)\n",
        "def d_tf_elu(x): return tf.cast(tf.greater(x,0),tf.float32)  + (tf_elu(tf.cast(tf.less_equal(x,0),tf.float32) * x) + 1.0)\n",
        "\n",
        "def tf_tanh(x): return tf.nn.tanh(x)\n",
        "def d_tf_tanh(x): return 1 - tf_tanh(x) ** 2\n",
        "\n",
        "def tf_sigmoid(x): return tf.nn.sigmoid(x) \n",
        "def d_tf_sigmoid(x): return tf_sigmoid(x) * (1.0-tf_sigmoid(x))\n",
        "\n",
        "def tf_atan(x): return tf.atan(x)\n",
        "def d_tf_atan(x): return 1.0/(1.0 + x**2)\n",
        "\n",
        "def tf_iden(x): return x\n",
        "def d_tf_iden(x): return 1.0\n",
        "\n",
        "def tf_softmax(x): return tf.nn.softmax(x)\n",
        "\n",
        "# code from: https://github.com/tensorflow/tensorflow/issues/8246\n",
        "def tf_repeat(tensor, repeats):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "\n",
        "    input: A Tensor. 1-D or higher.\n",
        "    repeats: A list. Number of repeat for each dimension, length must be the same as the number of dimensions in input\n",
        "\n",
        "    Returns:\n",
        "    \n",
        "    A Tensor. Has the same type as input. Has the shape of tensor.shape * repeats\n",
        "    \"\"\"\n",
        "    expanded_tensor = tf.expand_dims(tensor, -1)\n",
        "    multiples = [1] + repeats\n",
        "    tiled_tensor = tf.tile(expanded_tensor, multiples = multiples)\n",
        "    repeated_tesnor = tf.reshape(tiled_tensor, tf.shape(tensor) * repeats)\n",
        "    return repeated_tesnor\n",
        "\n",
        "# ================= VIZ =================\n",
        "# Def: Simple funciton to view the histogram of weights\n",
        "def show_hist_of_weigt(all_weight_list,status='before'):\n",
        "    fig = plt.figure()\n",
        "    weight_index = 0\n",
        "\n",
        "    for i in range(1,1+int(len(all_weight_list)//3)):\n",
        "        ax = fig.add_subplot(1,4,i)\n",
        "        ax.grid(False)\n",
        "        temp_weight_list = all_weight_list[weight_index:weight_index+3]\n",
        "        for temp_index in range(len(temp_weight_list)):\n",
        "            current_flat = temp_weight_list[temp_index].flatten()\n",
        "            ax.hist(current_flat,histtype='step',bins='auto',label=str(temp_index+weight_index))\n",
        "            ax.legend()\n",
        "        ax.set_title('From Layer : '+str(weight_index+1)+' to '+str(weight_index+3))\n",
        "        weight_index = weight_index + 3\n",
        "    plt.savefig('viz/weights_'+str(status)+\"_training.png\")\n",
        "    plt.close('all')\n",
        "\n",
        "# Def: Simple function to show 9 image with different channels\n",
        "def show_9_images(image,layer_num,image_num,channel_increase=3,alpha=None,gt=None,predict=None):\n",
        "    image = (image-image.min())/(image.max()-image.min())\n",
        "    fig = plt.figure()\n",
        "    color_channel = 0\n",
        "    limit = 10\n",
        "    if alpha: limit = len(gt)\n",
        "    for i in range(1,limit):\n",
        "        ax = fig.add_subplot(3,3,i)\n",
        "        ax.grid(False)\n",
        "        ax.set_xticks([])\n",
        "        ax.set_yticks([])\n",
        "        if alpha:\n",
        "            ax.set_title(\"GT: \"+str(gt[i-1])+\" Predict: \"+str(predict[i-1]))\n",
        "        else:\n",
        "            ax.set_title(\"Channel : \" + str(color_channel) + \" : \" + str(color_channel+channel_increase-1))\n",
        "        ax.imshow(np.squeeze(image[:,:,color_channel:color_channel+channel_increase]))\n",
        "        color_channel = color_channel + channel_increase\n",
        "    \n",
        "    if alpha:\n",
        "        plt.savefig('viz/z_'+str(alpha) + \"_alpha_image.png\")\n",
        "    else:\n",
        "        plt.savefig('viz/'+str(layer_num) + \"_layer_\"+str(image_num)+\"_image.png\")\n",
        "    plt.close('all')\n",
        "# ================= VIZ =================\n",
        "\n",
        "# ================= LAYER CLASSES =================\n",
        "class CNN():\n",
        "    \n",
        "    def __init__(self,k,inc,out,act=tf_elu,d_act=d_tf_elu):\n",
        "        self.w = tf.Variable(tf.random_normal([k,k,inc,out],stddev=0.05))\n",
        "        self.m,self.v_prev = tf.Variable(tf.zeros_like(self.w)),tf.Variable(tf.zeros_like(self.w))\n",
        "        self.act,self.d_act = act,d_act\n",
        "\n",
        "    def getw(self): return self.w\n",
        "\n",
        "    def feedforward(self,input,stride=1,padding='SAME'):\n",
        "        self.input  = input\n",
        "        self.layer  = tf.nn.conv2d(input,self.w,strides=[1,stride,stride,1],padding=padding) \n",
        "        self.layerA = self.act(self.layer)\n",
        "        return self.layerA \n",
        "\n",
        "    def backprop(self,gradient,stride=1,padding='SAME'):\n",
        "        grad_part_1 = gradient \n",
        "        grad_part_2 = self.d_act(self.layer) \n",
        "        grad_part_3 = self.input\n",
        "\n",
        "        grad_middle = grad_part_1 * grad_part_2\n",
        "\n",
        "        grad = tf.nn.conv2d_backprop_filter(input = grad_part_3,filter_sizes = self.w.shape,out_backprop = grad_middle,\n",
        "            strides=[1,stride,stride,1],padding=padding\n",
        "        )\n",
        "\n",
        "        grad_pass = tf.nn.conv2d_backprop_input(input_sizes = [batch_size] + list(grad_part_3.shape[1:]),filter= self.w,out_backprop = grad_middle,\n",
        "            strides=[1,stride,stride,1],padding=padding\n",
        "        )\n",
        "\n",
        "        update_w = []\n",
        "        update_w.append(tf.assign( self.m,self.m*beta1 + (1-beta1) * (grad)   ))\n",
        "        update_w.append(tf.assign( self.v_prev,self.v_prev*beta2 + (1-beta2) * (grad ** 2)   ))\n",
        "        m_hat = self.m / (1-beta1)\n",
        "        v_hat = self.v_prev / (1-beta2)\n",
        "        adam_middel = learning_rate/(tf.sqrt(v_hat) + adam_e)\n",
        "        update_w.append(tf.assign(self.w,tf.subtract(self.w,tf.multiply(adam_middel,m_hat)  )))         \n",
        "\n",
        "        return grad_pass,update_w \n",
        "\n",
        "class CNN_Trans():\n",
        "    \n",
        "    def __init__(self,k,inc,out,act=tf_elu,d_act=d_tf_elu):\n",
        "        self.w = tf.Variable(tf.random_normal([k,k,inc,out],stddev=0.05))\n",
        "        self.m,self.v_prev = tf.Variable(tf.zeros_like(self.w)),tf.Variable(tf.zeros_like(self.w))\n",
        "        self.act,self.d_act = act,d_act\n",
        "\n",
        "    def getw(self): return self.w\n",
        "\n",
        "    def feedforward(self,input,stride=1,padding='SAME'):\n",
        "        self.input  = input\n",
        "        output_shape2 = self.input.shape[2].value * stride\n",
        "        self.layer  = tf.nn.conv2d_transpose(\n",
        "            input,self.w,output_shape=[batch_size,output_shape2,output_shape2,self.w.shape[2].value],\n",
        "            strides=[1,stride,stride,1],padding=padding) \n",
        "        self.layerA = self.act(self.layer)\n",
        "        return self.layerA \n",
        "\n",
        "    def backprop(self,gradient,stride=1,padding='SAME'):\n",
        "        grad_part_1 = gradient \n",
        "        grad_part_2 = self.d_act(self.layer) \n",
        "        grad_part_3 = self.input\n",
        "\n",
        "        grad_middle = grad_part_1 * grad_part_2\n",
        "\n",
        "        grad = tf.nn.conv2d_backprop_filter(input = grad_middle,\n",
        "            filter_sizes = self.w.shape,out_backprop = grad_part_3,\n",
        "            strides=[1,stride,stride,1],padding=padding\n",
        "        )\n",
        "\n",
        "        grad_pass = tf.nn.conv2d(\n",
        "            input=grad_middle,filter = self.w,strides=[1,stride,stride,1],padding=padding\n",
        "        )\n",
        "        \n",
        "        update_w = []\n",
        "        update_w.append(tf.assign( self.m,self.m*beta1 + (1-beta1) * (grad)   ))\n",
        "        update_w.append(tf.assign( self.v_prev,self.v_prev*beta2 + (1-beta2) * (grad ** 2)   ))\n",
        "        m_hat = self.m / (1-beta1)\n",
        "        v_hat = self.v_prev / (1-beta2)\n",
        "        adam_middel = learning_rate/(tf.sqrt(v_hat) + adam_e)\n",
        "        update_w.append(tf.assign(self.w,tf.subtract(self.w,tf.multiply(adam_middel,m_hat)  )))         \n",
        "\n",
        "        return grad_pass,update_w \n",
        "\n",
        "class FNN():\n",
        "    \n",
        "    def __init__(self,input_dim,hidden_dim,act,d_act):\n",
        "        self.w = tf.Variable(tf.random_normal([input_dim,hidden_dim], stddev=0.05))\n",
        "        self.m,self.v_prev = tf.Variable(tf.zeros_like(self.w)),tf.Variable(tf.zeros_like(self.w))\n",
        "        self.v_hat_prev = tf.Variable(tf.zeros_like(self.w))\n",
        "        self.act,self.d_act = act,d_act\n",
        "\n",
        "    def feedforward(self,input=None):\n",
        "        self.input = input\n",
        "        self.layer = tf.matmul(input,self.w)\n",
        "        self.layerA = self.act(self.layer)\n",
        "        return self.layerA\n",
        "\n",
        "    def backprop(self,gradient=None):\n",
        "        grad_part_1 = gradient \n",
        "        grad_part_2 = self.d_act(self.layer) \n",
        "        grad_part_3 = self.input\n",
        "\n",
        "        grad_middle = grad_part_1 * grad_part_2\n",
        "        grad = tf.matmul(tf.transpose(grad_part_3),grad_middle)\n",
        "        grad_pass = tf.matmul(tf.multiply(grad_part_1,grad_part_2),tf.transpose(self.w))\n",
        "\n",
        "        update_w = []\n",
        "        update_w.append(tf.assign( self.m,self.m*beta1 + (1-beta1) * (grad)   ))\n",
        "        update_w.append(tf.assign( self.v_prev,self.v_prev*beta2 + (1-beta2) * (grad ** 2)   ))\n",
        "        m_hat = self.m / (1-beta1)\n",
        "        v_hat = self.v_prev / (1-beta2)\n",
        "        adam_middel = learning_rate/(tf.sqrt(v_hat) + adam_e)\n",
        "        update_w.append(tf.assign(self.w,tf.subtract(self.w,tf.multiply(adam_middel,m_hat)  )))     \n",
        "\n",
        "        return grad_pass,update_w   \n",
        "# ================= LAYER CLASSES =================\n",
        "\n",
        "# data\n",
        "data_location =  \"../../Dataset/salObj/datasets/imgs/ft/\"\n",
        "train_data = []  # create an empty list\n",
        "for dirName, subdirList, fileList in sorted(os.walk(data_location)):\n",
        "    for filename in fileList:\n",
        "        if \".jpg\" in filename.lower() :\n",
        "            train_data.append(os.path.join(dirName,filename))\n",
        "\n",
        "data_location = \"../../Dataset/salObj/datasets/masks/ft/\"\n",
        "train_data_gt = []  # create an empty list\n",
        "for dirName, subdirList, fileList in sorted(os.walk(data_location)):\n",
        "    for filename in fileList:\n",
        "        if \".png\" in filename.lower() :\n",
        "            train_data_gt.append(os.path.join(dirName,filename))\n",
        "\n",
        "image_resize_px = 64\n",
        "train_images = np.zeros(shape=(1000,image_resize_px,image_resize_px,3))\n",
        "train_labels = np.zeros(shape=(1000,image_resize_px,image_resize_px,1))\n",
        "\n",
        "for file_index in range(len(train_images)):\n",
        "    train_images[file_index,:,:]   = imresize(imread(train_data[file_index],mode='RGB'),(image_resize_px,image_resize_px))\n",
        "    train_labels[file_index,:,:]   = np.expand_dims(\n",
        "        imresize(rgb2gray(imread(train_data_gt[file_index],mode='RGB')),(image_resize_px,image_resize_px)),3)\n",
        "    \n",
        "train_images[:,:,:,0]  = (train_images[:,:,:,0] - train_images[:,:,:,0].min(axis=0)) / (train_images[:,:,:,0].max(axis=0) - train_images[:,:,:,0].min(axis=0)+1e-10)\n",
        "train_images[:,:,:,1]  = (train_images[:,:,:,1] - train_images[:,:,:,1].min(axis=0)) / (train_images[:,:,:,1].max(axis=0) - train_images[:,:,:,1].min(axis=0)+1e-10)\n",
        "train_images[:,:,:,2]  = (train_images[:,:,:,2] - train_images[:,:,:,2].min(axis=0)) / (train_images[:,:,:,2].max(axis=0) - train_images[:,:,:,2].min(axis=0)+1e-10)\n",
        "\n",
        "train_labels[:,:,:,0]  = (train_labels[:,:,:,0] - train_labels[:,:,:,0].min(axis=0)) / (train_labels[:,:,:,0].max(axis=0) - train_labels[:,:,:,0].min(axis=0)+1e-10)\n",
        "# train_labels[:,:,:,1]  = (train_labels[:,:,:,1] - train_labels[:,:,:,1].min(axis=0)) / (train_labels[:,:,:,1].max(axis=0) - train_labels[:,:,:,1].min(axis=0)+1e-10)\n",
        "# train_labels[:,:,:,2]  = (train_labels[:,:,:,2] - train_labels[:,:,:,2].min(axis=0)) / (train_labels[:,:,:,2].max(axis=0) - train_labels[:,:,:,2].min(axis=0)+1e-10)\n",
        "\n",
        "# split the data \n",
        "train_batch = train_images[:950]\n",
        "train_label = train_labels[:950]\n",
        "test_batch = train_images[950:]\n",
        "test_label = train_labels[950:]\n",
        "\n",
        "# train_batch = train_images[:50]\n",
        "# train_label = train_labels[:50]\n",
        "# test_batch = train_images[50:60]\n",
        "# test_label = train_labels[50:60]\n",
        "\n",
        "# print out the data shape\n",
        "print(train_batch.shape)\n",
        "print(train_label.shape)\n",
        "print(test_batch.shape)\n",
        "print(test_label.shape)\n",
        "\n",
        "# hyper parameter 10000\n",
        "num_epoch = 101\n",
        "batch_size = 10\n",
        "print_size = 2\n",
        "\n",
        "learning_rate = 0.000003\n",
        "learnind_rate_decay = 0.0\n",
        "beta1,beta2,adam_e = 0.9,0.999,1e-8\n",
        "\n",
        "# define class here\n",
        "el1 = CNN(3,3,128)\n",
        "el2 = CNN(3,128,256)\n",
        "el3 = CNN(3,256,512)\n",
        "\n",
        "e_mean = CNN(1,512,512)\n",
        "e_var  = CNN(1,512,512)\n",
        "\n",
        "dl1 = CNN_Trans(3,256,512)\n",
        "dl2 = CNN_Trans(3,128,256)\n",
        "dl3 = CNN_Trans(3,64,128)\n",
        "final_cnn = CNN(3,64,1,tf_sigmoid,d_tf_sigmoid)\n",
        "\n",
        "# graph\n",
        "x = tf.placeholder(shape=[batch_size,64,64,3],dtype=tf.float32,name=\"input\")\n",
        "y = tf.placeholder(shape=[batch_size,64,64,1],dtype=tf.float32,name=\"output\")\n",
        "\n",
        "# encoder\n",
        "elayer1 = el1.feedforward(x,padding='SAME')\n",
        "elayer2_input = tf.nn.avg_pool(elayer1,ksize=[1,2,2,1],strides=[1,2,2,1],padding='VALID')\n",
        "elayer2 = el2.feedforward(elayer2_input,padding='SAME')\n",
        "elayer3_input = tf.nn.avg_pool(elayer2,ksize=[1,2,2,1],strides=[1,2,2,1],padding='VALID')\n",
        "elayer3 = el3.feedforward(elayer3_input,padding='SAME')\n",
        "\n",
        "elayer_mean =  e_mean.feedforward(elayer3)\n",
        "elayer_var  =  e_var.feedforward(elayer3)\n",
        "eps = tf.random_normal((elayer_var.shape), 0, 1, dtype=tf.float32)\n",
        "z = elayer_mean + tf.sqrt(tf.exp(elayer_var)) * eps\n",
        "\n",
        "# edcoder\n",
        "dlayer1 = dl1.feedforward(z,stride=1,padding='SAME')\n",
        "dlayer2 = dl2.feedforward(dlayer1,stride=2,padding='SAME')\n",
        "dlayer3 = dl3.feedforward(dlayer2,stride=2,padding='SAME')\n",
        "final_output = final_cnn.feedforward(dlayer3,padding='SAME')\n",
        "\n",
        "# calculate the loss\n",
        "cost = tf.reduce_mean(tf.square(final_output-y)) * 0.5\n",
        "auto_train = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)\n",
        "\n",
        "# sess\n",
        "with tf.Session() as sess:\n",
        "\n",
        "    sess.run(tf.global_variables_initializer())\n",
        "\n",
        "    train_cota,train_acca = 0,0\n",
        "    train_cot,train_acc = [],[]\n",
        "    \n",
        "    test_cota,test_acca = 0,0\n",
        "    test_cot,test_acc = [],[]\n",
        "\n",
        "    # start the training\n",
        "    for iter in range(num_epoch):\n",
        "\n",
        "        train_batch,train_label = shuffle(train_batch,train_label)\n",
        "        test_batch,test_label = shuffle(test_batch,test_label)\n",
        "\n",
        "        for batch_size_index in range(0,len(train_batch),batch_size):\n",
        "            current_batch = train_batch[batch_size_index:batch_size_index+batch_size]\n",
        "            current_batch_labek = train_label[batch_size_index:batch_size_index+batch_size]\n",
        "            sess_result = sess.run([cost,auto_train],feed_dict={x:current_batch,y:current_batch_labek})\n",
        "            print(\"Current Iter : \",iter ,\" current batch: \",batch_size_index, ' Current cost: ', sess_result[0],end='\\r')\n",
        "            train_cota = train_cota + sess_result[0]\n",
        "\n",
        "        # For every print iteration save changes\n",
        "        if iter % print_size==0:\n",
        "            print(\"\\n--------------\")\n",
        "            print('Train Current cost: ', train_cota/(len(train_batch)/(batch_size)),end='\\n')\n",
        "            print(\"----------\")\n",
        "\n",
        "            # Get one example from training batch\n",
        "            test_example = train_batch[:batch_size,:,:,:]\n",
        "            test_example_gt = train_label[:batch_size,:,:,:]\n",
        "            sess_results = sess.run([final_output],feed_dict={x:test_example})\n",
        "            sess_results = sess_results[0][0,:,:,:]\n",
        "            test_example = test_example[0,:,:,:]\n",
        "            test_example_gt = test_example_gt[0,:,:,:]\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(test_example))\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image')\n",
        "            plt.savefig('train_change_train/'+str(iter)+\"a_Original_Image.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(test_example_gt),cmap='gray')\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image Mask')\n",
        "            plt.savefig('train_change_train/'+str(iter)+\"b_Original_Image_mask.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(test_example_gt*test_example)\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image Mask')\n",
        "            plt.savefig('train_change_train/'+str(iter)+\"c_Original_Image_overlay.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "  \n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(sess_results).astype(np.float32),cmap='gray')\n",
        "            plt.axis('off')\n",
        "            plt.title('Generated Mask')\n",
        "            plt.savefig('train_change_train/'+str(iter)+\"d_Generated_Mask.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(sess_results.astype(np.float32)*test_example)\n",
        "            plt.axis('off')\n",
        "            plt.title('Generated Mask')\n",
        "            plt.savefig('train_change_train/'+str(iter)+\"e_Generated_Mask_overlay.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            # Get one Example from test batch\n",
        "            test_example = test_batch[:batch_size,:,:,:]\n",
        "            test_example_gt = test_label[:batch_size,:,:,:]\n",
        "            sess_results = sess.run([final_output],feed_dict={x:test_example})\n",
        "            sess_results = sess_results[0][0,:,:,:]\n",
        "            test_example = test_example[0,:,:,:]\n",
        "            test_example_gt = test_example_gt[0,:,:,:]\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(test_example))\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image')\n",
        "            plt.savefig('train_change_test/'+str(iter)+\"a_Original_Image.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(test_example_gt),cmap='gray')\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image Mask')\n",
        "            plt.savefig('train_change_test/'+str(iter)+\"b_Original_Image_mask.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(test_example_gt*test_example)\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image Mask')\n",
        "            plt.savefig('train_change_test/'+str(iter)+\"c_Original_Image_overlay.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "  \n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(sess_results).astype(np.float32),cmap='gray')\n",
        "            plt.axis('off')\n",
        "            plt.title('Generated Mask')\n",
        "            plt.savefig('train_change_test/'+str(iter)+\"d_Generated_Mask.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(sess_results.astype(np.float32)*test_example)\n",
        "            plt.axis('off')\n",
        "            plt.title('Generated Mask')\n",
        "            plt.savefig('train_change_test/'+str(iter)+\"e_Generated_Mask_overlay.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "        # sort the training error\n",
        "        train_cot.append(train_cota/(len(train_batch)/(batch_size)))\n",
        "        train_cota,train_acca = 0,0\n",
        "\n",
        "    # Since all of the trainig is done create final test\n",
        "    for batch_size_index in range(0,len(test_batch),batch_size):\n",
        "        current_batch = test_batch[batch_size_index:batch_size_index+batch_size]\n",
        "        current_batch_labek = test_label[batch_size_index:batch_size_index+batch_size]\n",
        "        sess_result = sess.run(final_output,feed_dict={x:current_batch,y:current_batch_labek})\n",
        "\n",
        "        for sess_index in range(len(sess_result)):\n",
        "            sess_results = sess_result[sess_index]\n",
        "            test_example = test_batch[sess_index,:,:,:]\n",
        "            test_example_gt = test_label[sess_index,:,:,:]\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(test_example))\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image')\n",
        "            plt.savefig('viz/'+str(batch_size_index)+str(sess_index)+\"a_Original_Image.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(test_example_gt),cmap='gray')\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image Mask')\n",
        "            plt.savefig('viz/'+str(batch_size_index)+str(sess_index)+\"b_Original_Image_mask.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(test_example_gt*test_example)\n",
        "            plt.axis('off')\n",
        "            plt.title('Original Image Mask')\n",
        "            plt.savefig('viz/'+str(batch_size_index)+str(sess_index)+\"c_Original_Image_overlay.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "  \n",
        "            plt.figure()\n",
        "            plt.imshow(np.squeeze(sess_results).astype(np.float32),cmap='gray')\n",
        "            plt.axis('off')\n",
        "            plt.title('Generated Mask')\n",
        "            plt.savefig('viz/'+str(batch_size_index)+str(sess_index)+\"d_Generated_Mask.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(sess_results.astype(np.float32)*test_example)\n",
        "            plt.axis('off')\n",
        "            plt.title('Generated Mask')\n",
        "            plt.savefig('viz/'+str(batch_size_index)+str(sess_index)+\"e_Generated_Mask_overlay.png\",bbox_inches='tight')\n",
        "            plt.close('all')\n",
        "\n",
        "    # Normalize the cost of the training\n",
        "    train_cot = (train_cot-min(train_cot) ) / (max(train_cot)-min(train_cot))\n",
        "\n",
        "    # plot the training and testing graph\n",
        "    plt.figure()\n",
        "    plt.plot(range(len(train_acc)),train_acc,color='red',label='acc ovt')\n",
        "    plt.plot(range(len(train_cot)),train_cot,color='green',label='cost ovt')\n",
        "    plt.legend()\n",
        "    plt.title(\"Train Average Accuracy/Cost Over Time\")\n",
        "    plt.savefig(\"viz/z_Case Train.png\")\n",
        "    plt.close('all')\n",
        "\n",
        "# -- end code --"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}