{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Load Libraries"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":" %matplotlib inline\nimport os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt, zipfile\nimport matplotlib.image as mpimg\n\nimport tensorflow as tf\n\nimport os\nimport sys\nfrom tqdm import tqdm, tqdm_notebook\nimport shutil\n\n# import xml.etree.ElementTree as ET\nimport time # time the execution of codeblocks\n# from IPython.display import FileLink, FileLinks","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Path and list of directories"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"os.listdir('../input/')\npath = ('../input/')\nimg_output = path + 'img_output'\n\ntry:\n    if os.path.exists(img_output):\n        shutil.rmtree(img_output)\n    if not os.path.exists(img_output):\n        os.mkdir(img_output)\nexcept:\n    print('Unable to create the directory' + \"'\"+ img_output + \"'\")\n    print('Permission Denied!')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# List of dog breeds from annotation\ndog_breeds = os.listdir(path + 'annotation/Annotation')\nprint('total number of breeds %d'%len(dog_breeds))\n# pd.Series(dog_breeds).value_counts()\ndog_breeds = pd.DataFrame({'Annotation': dog_breeds})\ndog_breeds.loc[:,'Breed_Code'] = dog_breeds['Annotation'].apply(lambda x: x.split('-')[0])\ndog_breeds.loc[:,'Breed_Name'] = dog_breeds['Annotation'].apply(lambda x: x.split('-')[1])\ndog_breeds.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dog_img = os.listdir( path + 'all-dogs/all-dogs')\ndog_img = pd.DataFrame({'Image_Name': dog_img})\ndog_img.loc[:,'Breed_Code'] = dog_img['Image_Name'].apply(lambda x: x.split('_')[0])\ndog_img.loc[:,'Image_Num'] = dog_img['Image_Name'].apply(lambda x: x.split('_')[1]).apply(lambda x: x.split('.')[0])\ndog_img.loc[:,'Img_Path'] = dog_img['Image_Name'].apply(lambda x: path + 'all-dogs/all-dogs/' + x)\ndog_img.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dog Image Visualization"},{"metadata":{"trusted":true},"cell_type":"code","source":"\nimport random\nfrom PIL import Image\ndef plot_random_image(img_input, path, n_img = None, rows = None, img_size = None):\n    img_rand = random.sample(set(img_input), n_img)\n    cols = np.floor(len(img_rand)/rows)\n    plt.figure(figsize = (12,10))\n    for num, x in enumerate(img_rand):\n        img = Image.open(x)\n        img = img.resize((img_size, img_size), Image.ANTIALIAS)\n        plt.subplot(rows, cols, num + 1)\n        plt.axis('off')\n        plt.imshow(img)\n        plt.tight_layout()\n#         plt.show();\n#         if not os.path.exists(path + 'sample_dog_img/'):\n#             os.makedirs(path + 'sample_dog_img/')\n#         img.save(path + 'sample_dog_img/' + 'dog_' + str(num) + '.jpg')\n\nplot_random_image(dog_img.Img_Path.values, path, n_img = 10, rows = 2, img_size = 256)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Function to read image, resize and output image in sample batch"},{"metadata":{"trusted":true},"cell_type":"code","source":"###############################################################################\n# Function to read image, resize and output image in sample batch\n###############################################################################\n\n# def image_batch_process(image_path = dog_img['Img_Path'].values):\ndef image_batch_process():\n    img_path_tf = tf.convert_to_tensor(dog_img['Img_Path'].values, dtype = tf.string)   \n    img_path_tf = tf.train.slice_input_producer([img_path_tf])\n    raw_img = tf.read_file(img_path_tf[0])\n\n    img = tf.image.decode_jpeg(raw_img, channels = CHANNEL)\n\n    # img = tf.image.random_flip_left_right(img)\n    # img = tf.image.random_brightness(img, max_delta = 0.1)\n    # img = tf.image.random_contrast(img, lower = 0.9, upper = 1.1)\n\n    img = tf.image.resize_images(img, size = [HEIGHT, WIDTH])\n    img.set_shape([HEIGHT, WIDTH, CHANNEL])\n\n    img = tf.cast(img, tf.float32)\n    img = img / 255.0\n    \n    min_after_dequeue = 10000\n    capacity = min_after_dequeue + (THREADS + 1) * BATCH_SIZE\n\n    img_batch = tf.train.shuffle_batch([img], batch_size = BATCH_SIZE, \n                    num_threads = 10, capacity = capacity,\n                    min_after_dequeue = min_after_dequeue)\n    n_img = len(dog_img['Img_Path'].values)\n    return img_batch, n_img","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Function to train the Generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"###############################################################################\n# Function to train the Generator\n###############################################################################\n\n# Generator Function\ndef generator(input, RANDOM_DIM, is_train):\n    output_dim = CHANNEL  # RGB image\n    with tf.variable_scope('gen', reuse = tf.AUTO_REUSE):\n        \n        # Layer 1 = Input * W1 + b1 --> batch norm --> lrelu\n        gen_w1 = tf.get_variable('gen_w1', shape = [RANDOM_DIM, 4 * 4 * 256], dtype = tf.float32,\\\n                             initializer = tf.truncated_normal_initializer(stddev = 0.2))\n        gen_b1 = tf.get_variable('gen_b1', shape = [256 * 4 * 4], dtype = tf.float32,\\\n                             initializer = tf.constant_initializer(0.001))\n        layer_1 = tf.add(tf.matmul(input, gen_w1), gen_b1, name = 'layer_1')\n        \n        # Add Convolution layer1 --> batch norm --->relu\n        gen_conv1 = tf.reshape(layer_1, shape=[-1, 4, 4, 256], name = 'conv1')\n        gen_act1 = tf.nn.relu(gen_conv1, name = 'gen_act1')\n\n        # Add Convolution Layer2 --> batch norm --> relu\n        gen_conv2 = tf.layers.conv2d_transpose(gen_act1, 128, kernel_size = [5, 5], strides=[2, 2], padding=\"same\",\\\n                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\\\n                                           name='gen_conv2')        \n        gen_act2 = tf.nn.relu(gen_conv2, name = 'gen_act2')\n\n        # Add Convolution Layer3 --> batch norm --> relu\n        gen_conv3 = tf.layers.conv2d_transpose(gen_act2, 64, kernel_size=[5, 5], strides=[2, 2], padding=\"same\",\\\n                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\\\n                                           name = 'gen_conv3')\n        gen_act3 = tf.nn.relu(gen_conv3, name='gen_act3')\n        \n        # Add Convolution Layer4 --> batch norm --> relu\n        gen_conv4 = tf.layers.conv2d_transpose(gen_act3, 32, kernel_size=[5, 5], strides=[2, 2], padding=\"same\",\\\n                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\\\n                                           name='gen_conv4')\n        gen_act4 = tf.nn.relu(gen_conv4, name = 'gen_act4')\n    \n        # Add Convolution Layer5 --> tanh\n        gen_conv5 = tf.layers.conv2d_transpose(gen_act4, output_dim, kernel_size=[5, 5], strides=[2, 2], padding=\"same\",\\\n                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\\\n                                           name='gen_conv5')\n        gen_act5 = tf.nn.tanh(gen_conv5, name = 'gen_act5')\n        print('Gen Activation 5 shape:', gen_act5.shape)\n        return gen_act5","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Function to train the Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"###############################################################################\n# Function to train the Discriminator\n###############################################################################\n\n# Discriminator Function\ndef discriminator(input, is_train, reuse = False):\n    print('Inside Discriminator...')\n    with tf.variable_scope('dis', reuse = tf.AUTO_REUSE):\n#     with tf.variable_scope('dis') as scope:\n#         if reuse:\n#             scope.reuse_variables()\n\n        # Add Convolution layer1 --> batch norm --->relu\n        dis_conv1 = tf.layers.conv2d(input, 64, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\\\n                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\\\n                                           name='dis_conv1')        \n        dis_act1 = tf.nn.leaky_relu(dis_conv1, name='dis_act1')\n\n        # Add Convolution Layer2 --> batch norm --> relu\n        dis_conv2 = tf.layers.conv2d(dis_act1, 128, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\\\n                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\\\n                                           name='dis_conv2')\n        dis_act2 = tf.nn.leaky_relu(dis_conv2, name='dis_act2')\n\n        # Add Convolution Layer3 --> batch norm --> relu\n        dis_conv3 = tf.layers.conv2d(dis_act2, 256, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\\\n                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\\\n                                           name='dis_conv3')\n        dis_act3 = tf.nn.leaky_relu(dis_conv3, name='dis_act3')\n\n        # Here we are re-shaping the tensor to feed into Wasserstein function\n        # to check and compete between - Fake and Original image by generator and discriminator\n        dim = np.int(np.prod(dis_act3.get_shape()[1:]))\n        conv_vect = tf.reshape(dis_act3, shape = [-1, dim], name = 'conv_vect')\n\n        # Weight Initialization\n        dis_w2 = tf.get_variable('dis_w2', shape = [conv_vect.shape[-1], 1], dtype = tf.float32,\\\n                             initializer=tf.truncated_normal_initializer(stddev = 0.02))\n        dis_b2 = tf.get_variable('dis_b2', shape = [1], dtype = tf.float32,\\\n                             initializer = tf.constant_initializer(0.001))\n\n        # Wasserstein GAN - wgan to caculate Wasserstein Distance to check generator \n        dis_wgan = tf.add(tf.matmul(conv_vect, dis_w2), dis_b2, name = 'dis_wgan')\n        print('Shape of Linear Wgan:', dis_wgan.shape)\n\n        return dis_wgan","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Function to train the Discriminator/Generator Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"###############################################################################\n# Function to train the Discriminator/Generator Model\n###############################################################################\n\n# Main Train function\ndef train():\n\n    with tf.variable_scope('inputs', reuse = tf.AUTO_REUSE):\n        real_data = tf.placeholder(tf.float32, shape = [None, HEIGHT, WIDTH, CHANNEL], name = 'real_image')\n        random_input = tf.placeholder(tf.float32, shape=[None, RANDOM_DIM], name = 'rand_input')\n        is_train = tf.placeholder(tf.bool, name = 'is_train')\n    \n    # Generator function call\n    fake_data = generator(random_input, RANDOM_DIM, is_train = True)\n    \n    # Training two discriminators to discriminate both real and fake data\n    real_result = discriminator(real_data, is_train)\n    fake_result = discriminator(fake_data, is_train, reuse = True)\n    \n    # Loss functions of Discriminator and Generator\n    # d_loss --> will optimize discriminator\n    d_loss = tf.reduce_mean(fake_result) - tf.reduce_mean(real_result)\n    \n    # g_loss --> will optimize generator\n    g_loss = -tf.reduce_mean(fake_result)\n            \n    # Getting the list of trainable variables/weights from generator and descriminator\n    t_vars = tf.trainable_variables()\n\n    # List of Trainable Variables\n\n    d_vars = [var for var in t_vars if 'dis' in var.name]\n    g_vars = [var for var in t_vars if 'gen' in var.name]\n#    print('Trainable variables in Dis: ', d_vars)\n#    print('Trainable variables in Gen: ', g_vars)\n    \n# running optimization on trainble variables\n    with tf.variable_scope('optimizer', reuse = tf.AUTO_REUSE):\n#         trainer_d = tf.train.RMSPropOptimizer(learning_rate = LEARNING_RATE, \\\n#                        name = 'dis_rmsprop').minimize(d_loss, var_list = d_vars)\n        \n        trainer_d = tf.train.AdamOptimizer(learning_rate = LEARNING_RATE, \\\n#                     beta1 = 0.5, beta2 = 0.9,\\\n                    name = 'dis_adam').minimize(d_loss, var_list = d_vars)\n        \n#         trainer_g = tf.train.RMSPropOptimizer(learning_rate = LEARNING_RATE, \\\n#                        name = 'gen_rmsprop').minimize(g_loss, var_list = g_vars)\n        \n        trainer_g = tf.train.AdamOptimizer(learning_rate = LEARNING_RATE, \\\n                     #   beta1 = 0.5, beta2 = 0.9,\\\n                        name = 'gen_adam').minimize(g_loss, var_list = g_vars)\n    \n    # clip discriminator weights. Restricting the weights of discriminator\n    d_clip = [v.assign(tf.clip_by_value(v, -0.01, 0.01)) for v in d_vars]\n\n    # Fetch image data from image process function\n    img_batch, samples_num = image_batch_process()\n    batch_num = int(samples_num / BATCH_SIZE)\n    \n    # TF Session start\n    sess = tf.Session()\n    saver = tf.train.Saver()\n    sess.run(tf.global_variables_initializer())\n    sess.run(tf.local_variables_initializer())\n    \n    coord = tf.train.Coordinator()\n    threads = tf.train.start_queue_runners(sess=sess, coord=coord)\n    \n    print('\\ntotal training sample:%d' % samples_num)\n    print('\\nbatch size: %d, batch num per epoch: %d, epoch num: %d' \\\n          % (BATCH_SIZE, batch_num, EPOCH))\n    print('\\nstart training...')\n    \n    # List to store avg generator and discriminator loss\n    avg_gen_loss = []\n    avg_dis_loss = []\n    for i in range(EPOCH):\n        print(\"\\nRunning epoch {}/{}...\".format(i, EPOCH))\n        \n        gen_loss = []\n        dis_loss = []\n        for j in range(batch_num):\n#             print('Training on Batch number %d'%j)\n            d_iters = 3\n            g_iters = 1\n\n            train_noise = np.random.uniform(-1.0, 1.0, size = [BATCH_SIZE, \\\n                                            RANDOM_DIM]).astype(np.float32)\n\n#            train_noise = np.random.normal(0.0, 1.0, size = [BATCH_SIZE, \\\n#                                            RANDOM_DIM]).astype(np.float32)\n\n    # Training a Discriminator 3 times for a single Generator\n            for k in range(d_iters):\n#                 print('Training Discriminator: ', k)\n                img_train = sess.run(img_batch)\n#                 print('Image Training Batch Size: ', img_train.shape)\n                \n                #wgan clip weights\n                sess.run(d_clip)\n                    \n                    # Update the discriminator\n                _, disLoss = sess.run([trainer_d, d_loss],\\\n                         feed_dict = {random_input: train_noise, real_data: img_train, \\\n                                                   is_train: True})\n#             print('Discriminator Loss:', disLoss)\n            dis_loss.append(disLoss)\n\n            # Update the generator\n            for k in range(g_iters):\n#                 print('Training Generator: ',k)\n#                 print('Generator Noise input shape: ', train_noise.shape)\n               # train_noise = np.random.uniform(-1.0, 1.0, size=[BATCH_SIZE, RANDOM_DIM]).astype(np.float32)\n                _, genLoss = sess.run([trainer_g, g_loss],\\\n                            feed_dict = {random_input: train_noise, \\\n                                                   is_train: True})\n#             print('Generator Loss: ', genLoss)\n            gen_loss.append(genLoss)\n        \n        # Caculating avg Generator and Discriminator loss\n        avg_gen_loss.append(np.mean(gen_loss))\n        avg_dis_loss.append(np.mean(dis_loss))\n\n        print('Epochs:[%d/%d], avg_dis_loss:%f, avg_gen_loss:%f' % (i, EPOCH, np.mean(dis_loss), \\\n              np.mean(gen_loss)))\n            \n        # save check point every 500 epoch\n#         if i%2 == 0:\n#             try:\n#                 if not os.path.exists(dog_img_path +  version +'model/'):\n#                     os.makedirs(path +  version + 'model/')\n#                 saver.save(sess, path +  version + 'model/'+ str(i) + '_dog_img_gen.ckpt')\n#             except:\n#                 print('Unable to save the trained session. permission denied!')\n        final_epoch = i + 1\n        if final_epoch%EPOCH == 0:   \n            # save Generated Samples\n#             try:\n#                 if not os.path.exists(path +  version + 'Dog_Gen_Images/'):\n#                     os.makedirs(path +  version + 'Dog_Gen_Images/')\n#             except:\n#                 print('Unable to create path for saving image!')\n            \n            sample_noise = np.random.uniform(-1.0, 1.0, size = [SAMPLE_SIZE, RANDOM_DIM]).astype(np.float32)\n            sample_pred = sess.run(fake_data, feed_dict = {random_input: sample_noise, is_train: False})\n            \n            print('Sample Generated Shape: ', sample_pred.shape)\n            print('Showing some random image generated...')\n            \n            xxx = 0\n            for smp_img in sample_pred:\n                imgs = smp_img * 255.0\n#                 imgs = imgs.numpy()\n#                 print('SMP_IMG: ', smp_img[0])\n#                 imgs = Image.fromarray(smp_img.astype('uint8').reshape((64,64)), 'RGB')\n                imgs = Image.fromarray(smp_img.astype('uint8'), 'RGB')\n#                 sample_images(imgs, rows = HEIGHT, cols=WIDTH)\n                plt.figure(figsize = (8,5))\n                plt.title('Reconstructed')\n                plt.imshow(imgs)\n                plt.axis('off')\n                xxx +=1\n        \n        final_epoch = i + 1\n        if final_epoch%EPOCH == 0:\n            # save Generated Samples\n            sample_noise = np.random.uniform(-1.0, 1.0, size = [10000, RANDOM_DIM]).astype(np.float32)\n            sample_pred = sess.run(fake_data, feed_dict = {random_input: sample_noise, is_train: False})\n            \n#             print('Sample Generated Shape: ', sample_pred.shape)\n            \n            # SAVE TO ZIP FILE NAMED IMAGES.ZIP\n            z = zipfile.PyZipFile('images.zip', mode='w')\n            lop = 0\n            try:\n                for smp_img in sample_pred:\n                    imgs = smp_img * 255.0\n#                     imgs = imgs.numpy()\n                    imgs = Image.fromarray(imgs.astype('uint8').reshape((64,64,3)), 'RGB')\n#                     imgs.save(img_output + '/dog_' + str(lop).zfill(3) + '.jpg')\n                    f = str(lop).zfill(3) + '.png'\n                    tf.keras.preprocessing.image.save_img(f, imgs, scale=True)\n                    lop +=1\n#                 print('Generated Image in Pixels: ', np.asarray(imgs, dtype=\"int32\" ))\n                    z.write(f); os.remove(f)\n                z.close()\n                \n            except:\n                print('Unable to create and save image in directory!')\n            # Plot some random image\n#             plot_random_image(sample_pred, path, n_img = 10, rows = 2, img_size = HEIGHT)\n        \n    coord.request_stop()\n    coord.join(threads)\n    \n    print('Plotting Average Generator and Discriminator Loss...')\n    plt.figure(figsize = (12,10))\n    plt.plot(avg_gen_loss)\n    plt.plot(avg_dis_loss)\n    plt.title('Generator & Discriminator Loss by Epochs')\n    plt.ylabel('Loss')\n    plt.xlabel('Epochs')\n    plt.legend(['Generator Loss', 'Discriminator Loss'], loc = 'upper right')\n    plt.show()\n    \n    # Saving the complete trained session/model\n    \n    try:\n        if not os.path.exists(path + version + 'final_model/'):\n            os.makedirs(path + version +'final_model/')\n        saver.save(sess, path + version +'final_model/' + 'wgan_dog_img_gen.ckpt')\n        print('\\nFinal WGAN Model Stored path: ')\n    except:\n        print('Unable to save the final model!')\n\n    print('\\nFinal Generator Loss: ', np.mean(avg_gen_loss))\n    print('Final Discriminator Loss: ', np.mean(avg_dis_loss))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Main Function to call all the above functions and train the model"},{"metadata":{"trusted":true},"cell_type":"code","source":"###############################################################################\n# Main Function to call all the above functions\n###############################################################################\n\nif __name__ == \"__main__\":\n    # Major Hyper-parameters\n    THREADS = 4\n    BATCH_SIZE = 128\n    EPOCH = 300\n    RANDOM_DIM = 128\n    LEARNING_RATE = 0.0003\n    SAMPLE_SIZE = 128\n    version = 'GAN_DOG_IMG_V1/'\n    HEIGHT, WIDTH, CHANNEL = 64, 64, 3\n    import time\n    t1 = time.time()\n    train()\n    print('total time taken in mins: ', (time.time()  - t1)/60)\n#    data_generator(nrows = 10000)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}