DAE.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DAE.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Aftaab99/DenoisingAutoencoder/blob/master/DAE.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "Mq_UucuWJdUQ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "D_GMWcrgKNj-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.chdir(\"/content/drive/My Drive/DenoisingAutoencoder/Data\")\n",
        "!ls"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "2ZT4PFeUKyg1",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.python.keras.layers import Convolution2D, Convolution2DTranspose, merge, Input\n",
        "from tensorflow.python.keras.models import Model\n",
        "# os.mkdir('Checkpoints')\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
        "tf.enable_eager_execution()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "qT6_a2HWLLhD",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "n_records = 0\n",
        "for record in tf.python_io.tf_record_iterator('train.tfrecords'):\n",
        "\tn_records += 1\n",
        "\n",
        "batch_size = 2048\n",
        "\n",
        "class DAE:\n",
        "\n",
        "\tdef __init__(self):\n",
        "\t\tinput_1 = Input(shape=(None, None, 3))\n",
        "\t\tconv_1 = Convolution2D(64, kernel_size=(3, 3), padding='same', activation='relu')(input_1)\n",
        "\t\tconv_2 = Convolution2D(64, kernel_size=(5, 5), padding='same', activation='relu')(conv_1)\n",
        "\t\tdconv_1 = Convolution2DTranspose(64, kernel_size=(3, 3), padding='same', activation='relu')(conv_2)\n",
        "\t\tmerge_1 = merge.maximum([dconv_1, conv_2])\n",
        "\t\tdconv_2 = Convolution2DTranspose(64, kernel_size=(3, 3), padding=\"same\", activation='relu')(merge_1)\n",
        "\t\tmerge_2 = merge.maximum([dconv_2, conv_1])\n",
        "\t\tconv3 = Convolution2D(3, (5, 5), padding=\"same\", activation='relu')(merge_2)\n",
        "\n",
        "\t\tself.model = Model(inputs=input_1, outputs=conv3)\n",
        "\t\tself.model.compile(optimizer='adam', loss='mean_squared_error', metrics=['acc'])\n",
        "\t\tself.model.summary()\n",
        "\n",
        "\tdef load_model(self, save_path):\n",
        "\t\tself.model.load_weights(save_path)\n",
        "\n",
        "\tdef save_model(self, save_path):\n",
        "\t\tself.model.save_weights(save_path)\n",
        "\n",
        "\tdef train(self, epochs, train_records_filename):\n",
        "\t\tx, y = self.input_fn(train_records_filename)\n",
        "\t\tself.model.fit(x.numpy(), y.numpy(), epochs=epochs, steps_per_epoch=n_records // batch_size)\n",
        "\n",
        "\tdef denoise_patch(self, image_patch):\n",
        "\t\timage_patch = image_patch[np.newaxis, ...]\n",
        "\t\toutput_t = self.model.predict(image_patch)\n",
        "\t\toutput_t = np.array(output_t)\n",
        "\t\toutput_t = np.clip(output_t, 0, 255)\n",
        "\t\treturn output_t\n",
        "\n",
        "\tdef denoise(self, image_array):\n",
        "\t\tdim = image_array.shape\n",
        "\t\timg_h = dim[0]\n",
        "\t\timg_w = dim[1]\n",
        "\t\td_image = image_array\n",
        "\t\t# print(img_h*img_w)\n",
        "\t\tif img_w * img_h < 400 * 400:\n",
        "\t\t\timage_array = image_array[np.newaxis, ...]\n",
        "\t\t\ta = np.clip(d.model.predict(image_array), 0, 255).astype('uint8')\n",
        "\t\t\ta = a.squeeze(0)\n",
        "\t\t\timg1 = Image.fromarray(a)\n",
        "\t\t\treturn img1\n",
        "\n",
        "\t\tfor y in range(0, img_w, 33):\n",
        "\t\t\tfor x in range(0, img_h, 33):\n",
        "\t\t\t\tpatch = image_array[x:x + 33, y:y + 33, :]\n",
        "\t\t\t\tprint(patch.shape)\n",
        "\t\t\t\tif patch.shape[0] == 33 and patch.shape[1] == 33:\n",
        "\t\t\t\t\tpatch = self.denoise_patch(patch)\n",
        "\t\t\t\t\td_image[x:x + 33, y:y + 33, :] = patch\n",
        "\n",
        "\n",
        "\t\t\t\telif patch.shape[0] < 33 and patch.shape[1] < 33:\n",
        "\t\t\t\t\tpatch = self.denoise_patch(patch)\n",
        "\t\t\t\t\td_image[x:, y:, :] = patch\n",
        "\n",
        "\n",
        "\t\t\t\telif patch.shape[1] < 33 and patch.shape[0] == 33:\n",
        "\t\t\t\t\tl = patch.shape[1]\n",
        "\t\t\t\t\tpatch = self.denoise_patch(patch)\n",
        "\t\t\t\t\td_image[x:x + 33, y:y + l, :] = patch\n",
        "\n",
        "\t\t\t\telif patch.shape[0] < 33 and patch.shape[1] == 33:\n",
        "\t\t\t\t\tl = patch.shape[0]\n",
        "\t\t\t\t\tpatch = self.denoise_patch(patch)\n",
        "\t\t\t\t\td_image[x:x + l, y:y + 33, :] = patch[0:l, :, :]\n",
        "\n",
        "\t\td_image = Image.fromarray(d_image.astype('uint8'))\n",
        "\t\treturn d_image\n",
        "\n",
        "\tdef parser(self, record):\n",
        "\t\tkeys_to_feature = {\n",
        "\t\t\t\"reference\": tf.FixedLenFeature([], tf.string),\n",
        "\t\t\t\"noisy\": tf.FixedLenFeature([], tf.string)\n",
        "\t\t}\n",
        "\t\tparsed = tf.parse_single_example(record, keys_to_feature)\n",
        "\t\ttarget_image = tf.decode_raw(parsed['reference'], tf.uint8)\n",
        "\t\ttarget_image = tf.cast(target_image, tf.float32)\n",
        "\n",
        "\t\ttarget_image = tf.reshape(target_image, shape=[33, 33, 3])\n",
        "\t\tnoisy_image = tf.decode_raw(parsed['noisy'], tf.uint8)\n",
        "\t\tnoisy_image = tf.cast(noisy_image, tf.float32)\n",
        "\t\tnoisy_image = tf.reshape(noisy_image, shape=[33, 33, 3])\n",
        "\t\treturn noisy_image, target_image\n",
        "\n",
        "\tdef input_fn(self, filename):\n",
        "\t\tdataset = tf.data.TFRecordDataset(filename)\n",
        "\t\tdataset = dataset.map(self.parser)\n",
        "\t\tdataset = dataset.repeat()\n",
        "\n",
        "\t\tdataset = dataset.batch(batch_size)\n",
        "\t\titerator = dataset.make_one_shot_iterator()\n",
        "\t\tnoisy_batch, target_batch = iterator.get_next()\n",
        "\t\treturn noisy_batch, target_batch"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "PlVFhnHrLXvC",
        "colab_type": "code",
        "outputId": "d270ac1b-9900-491e-9bbe-560af1a7cd1b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1395
        }
      },
      "cell_type": "code",
      "source": [
        "d = DAE()\n",
        "d.train(epochs=25, train_records_filename='train.tfrecords')\n",
        "d.save_model('model_keras_weights_final.hdf5')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "__________________________________________________________________________________________________\n",
            "Layer (type)                    Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            "input_4 (InputLayer)            (None, None, None, 3 0                                            \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_9 (Conv2D)               (None, None, None, 6 1792        input_4[0][0]                    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_10 (Conv2D)              (None, None, None, 6 102464      conv2d_9[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_transpose_6 (Conv2DTrans (None, None, None, 6 36928       conv2d_10[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "maximum_6 (Maximum)             (None, None, None, 6 0           conv2d_transpose_6[0][0]         \n",
            "                                                                 conv2d_10[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_transpose_7 (Conv2DTrans (None, None, None, 6 36928       maximum_6[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "maximum_7 (Maximum)             (None, None, None, 6 0           conv2d_transpose_7[0][0]         \n",
            "                                                                 conv2d_9[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_11 (Conv2D)              (None, None, None, 3 4803        maximum_7[0][0]                  \n",
            "==================================================================================================\n",
            "Total params: 182,915\n",
            "Trainable params: 182,915\n",
            "Non-trainable params: 0\n",
            "__________________________________________________________________________________________________\n",
            "Epoch 1/25\n",
            "104/104 [==============================] - 85s 821ms/step - loss: 381.4120 - acc: 0.7720\n",
            "Epoch 2/25\n",
            "104/104 [==============================] - 88s 846ms/step - loss: 90.2350 - acc: 0.8113\n",
            "Epoch 3/25\n",
            "104/104 [==============================] - 88s 848ms/step - loss: 81.8493 - acc: 0.8202\n",
            "Epoch 4/25\n",
            "104/104 [==============================] - 88s 850ms/step - loss: 78.2569 - acc: 0.8243\n",
            "Epoch 5/25\n",
            "104/104 [==============================] - 88s 849ms/step - loss: 74.3489 - acc: 0.8290\n",
            "Epoch 6/25\n",
            "104/104 [==============================] - 88s 850ms/step - loss: 72.9516 - acc: 0.8296\n",
            "Epoch 7/25\n",
            "104/104 [==============================] - 88s 850ms/step - loss: 70.4544 - acc: 0.8332\n",
            "Epoch 8/25\n",
            "104/104 [==============================] - 88s 850ms/step - loss: 70.9401 - acc: 0.8323\n",
            "Epoch 9/25\n",
            "104/104 [==============================] - 89s 851ms/step - loss: 66.7533 - acc: 0.8360\n",
            "Epoch 10/25\n",
            "104/104 [==============================] - 88s 851ms/step - loss: 65.8266 - acc: 0.8367\n",
            "Epoch 11/25\n",
            "104/104 [==============================] - 88s 850ms/step - loss: 64.3367 - acc: 0.8377\n",
            "Epoch 12/25\n",
            "104/104 [==============================] - 88s 850ms/step - loss: 63.6573 - acc: 0.8377\n",
            "Epoch 13/25\n",
            "104/104 [==============================] - 88s 850ms/step - loss: 62.2400 - acc: 0.8379\n",
            "Epoch 14/25\n",
            "104/104 [==============================] - 88s 851ms/step - loss: 61.0746 - acc: 0.8397\n",
            "Epoch 15/25\n",
            "104/104 [==============================] - 89s 851ms/step - loss: 60.5733 - acc: 0.8398\n",
            "Epoch 16/25\n",
            "104/104 [==============================] - 89s 852ms/step - loss: 59.7295 - acc: 0.8402\n",
            "Epoch 17/25\n",
            "104/104 [==============================] - 88s 851ms/step - loss: 58.7280 - acc: 0.8414\n",
            "Epoch 18/25\n",
            "104/104 [==============================] - 88s 851ms/step - loss: 58.1715 - acc: 0.8416\n",
            "Epoch 19/25\n",
            "104/104 [==============================] - 88s 851ms/step - loss: 57.8916 - acc: 0.8410\n",
            "Epoch 20/25\n",
            "104/104 [==============================] - 89s 853ms/step - loss: 56.5707 - acc: 0.8429\n",
            "Epoch 21/25\n",
            "104/104 [==============================] - 89s 854ms/step - loss: 56.2133 - acc: 0.8436\n",
            "Epoch 22/25\n",
            "104/104 [==============================] - 89s 852ms/step - loss: 56.2928 - acc: 0.8436\n",
            "Epoch 23/25\n",
            "104/104 [==============================] - 89s 853ms/step - loss: 54.9715 - acc: 0.8442\n",
            "Epoch 24/25\n",
            "104/104 [==============================] - 89s 853ms/step - loss: 64.5708 - acc: 0.8390\n",
            "Epoch 25/25\n",
            "104/104 [==============================] - 88s 851ms/step - loss: 59.6806 - acc: 0.8397\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "4i_Yn8LC9uCd",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}