Single_Image_Dehazing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QJPU0tPrAAAi"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import PIL\n",
    "import glob\n",
    "import torch\n",
    "import scipy\n",
    "import random\n",
    "import shutil\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from math import exp\n",
    "import torch.nn as nn\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "from scipy import ndimage\n",
    "import torch.optim as optim\n",
    "from sklearn import metrics\n",
    "from torch.utils import data\n",
    "from sklearn import datasets\n",
    "from skimage.io import imsave\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn.functional as F3\n",
    "import torch.nn.functional as F9\n",
    "from sklearn import linear_model\n",
    "from skimage.feature import canny\n",
    "from torch.autograd import Variable\n",
    "import torchvision.models as models\n",
    "from collections import OrderedDict\n",
    "from torch.autograd import Variable\n",
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt1\n",
    "from torch.utils.data import DataLoader\n",
    "from sklearn.model_selection import KFold\n",
    "from IPython.display import display, Image\n",
    "from skimage.color import lab2rgb, rgb2lab\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.transforms.functional as F\n",
    "import torchvision.transforms.functional as F6\n",
    "import torchvision.transforms.functional as F7\n",
    "from sklearn.preprocessing import OneHotEncoder \n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import mean_squared_error,r2_score\n",
    "from skimage.color import rgb2lab, lab2rgb, rgb2gray, xyz2lab\n",
    "from sklearn.model_selection import cross_val_score, cross_val_predict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3OcaGLs7bOiE"
   },
   "source": [
    "Defining our generator (UNet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uPQd3FlmAAAk"
   },
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "\n",
    "    def __init__(self, in_channels=3, out_channels=3, init_features=32):\n",
    "        super(UNet, self).__init__()\n",
    "\n",
    "        features = init_features\n",
    "        self.encoder1 = UNet._block(in_channels, features, name=\"enc1\")\n",
    "        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.encoder2 = UNet._block(features, features * 2, name=\"enc2\")\n",
    "        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.encoder3 = UNet._block(features * 2, features * 4, name=\"enc3\")\n",
    "        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.encoder4 = UNet._block(features * 4, features * 8, name=\"enc4\")\n",
    "        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "        self.bottleneck = UNet._block(features * 8, features * 16, name=\"bottleneck\")\n",
    "\n",
    "        self.upconv4 = nn.ConvTranspose2d(\n",
    "            features * 16, features * 8, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder4 = UNet._block((features * 8) * 2, features * 8, name=\"dec4\")\n",
    "        self.upconv3 = nn.ConvTranspose2d(\n",
    "            features * 8, features * 4, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder3 = UNet._block((features * 4) * 2, features * 4, name=\"dec3\")\n",
    "        self.upconv2 = nn.ConvTranspose2d(\n",
    "            features * 4, features * 2, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder2 = UNet._block((features * 2) * 2, features * 2, name=\"dec2\")\n",
    "        self.upconv1 = nn.ConvTranspose2d(\n",
    "            features * 2, features, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder1 = UNet._block(features * 2, features, name=\"dec1\")\n",
    "\n",
    "        self.conv = nn.Conv2d(\n",
    "            in_channels=features, out_channels=out_channels, kernel_size=1\n",
    "        )\n",
    "\n",
    "        self.conv_1 = nn.Conv2d(\n",
    "            in_channels=128, out_channels=out_channels, kernel_size=1\n",
    "        )\n",
    "\n",
    "        self.pyramid_3 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=3, padding=1),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        self.pyramid_5 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=5, padding=2),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        self.pyramid_7 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=7, padding=3),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        self.pyramid_11 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=11, padding=5),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        self.pyramid_17 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=17, padding=8),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        self.pyramid_25 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=25, padding=12),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        self.pyramid_35 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=35, padding=17),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        self.pyramid_45 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=12, out_channels=16, kernel_size=45, padding=22),\n",
    "            nn.InstanceNorm2d(16, track_running_stats=False),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        #block1\n",
    "        enc1 = self.encoder1(x)\n",
    "        enc2 = self.encoder2(self.pool1(enc1))\n",
    "        enc3 = self.encoder3(self.pool2(enc2))\n",
    "        enc4 = self.encoder4(self.pool3(enc3))\n",
    "\n",
    "        bottleneck = self.bottleneck(self.pool4(enc4))\n",
    "\n",
    "        dec4 = self.upconv4(bottleneck)\n",
    "        dec4 = torch.cat((dec4, enc4), dim=1)\n",
    "        dec4 = self.decoder4(dec4)\n",
    "        dec3 = self.upconv3(dec4)\n",
    "        dec3 = torch.cat((dec3, enc3), dim=1)\n",
    "        dec3 = self.decoder3(dec3)\n",
    "        dec2 = self.upconv2(dec3)\n",
    "        dec2 = torch.cat((dec2, enc2), dim=1)\n",
    "        dec2 = self.decoder2(dec2)\n",
    "        dec1 = self.upconv1(dec2)\n",
    "        dec1 = torch.cat((dec1, enc1), dim=1)\n",
    "        dec1 = self.decoder1(dec1)\n",
    "        coonv1_1 = self.conv(dec1)\n",
    "\n",
    "        \n",
    "        #block2 \n",
    "        enc1_2 = self.encoder1(coonv1_1)\n",
    "        enc2_2 = self.encoder2(self.pool1(enc1_2))\n",
    "        enc3_2 = self.encoder3(self.pool2(enc2_2))\n",
    "        enc4_2 = self.encoder4(self.pool3(enc3_2))\n",
    "\n",
    "        bottleneck_2 = self.bottleneck(self.pool4(enc4_2))\n",
    "\n",
    "        dec4_2 = self.upconv4(bottleneck_2)\n",
    "        dec4_2 = torch.cat((dec4_2, enc4_2), dim=1)\n",
    "        dec4_2 = self.decoder4(dec4_2)\n",
    "        dec3_2 = self.upconv3(dec4_2)\n",
    "        dec3_2 = torch.cat((dec3_2, enc3_2), dim=1)\n",
    "        dec3_2 = self.decoder3(dec3_2)\n",
    "        dec2_2 = self.upconv2(dec3_2)\n",
    "        dec2_2 = torch.cat((dec2_2, enc2_2), dim=1)\n",
    "        dec2_2 = self.decoder2(dec2_2)\n",
    "        dec1_2 = self.upconv1(dec2_2)\n",
    "        dec1_2 = torch.cat((dec1_2, enc1_2), dim=1)\n",
    "        dec1_2 = self.decoder1(dec1_2)\n",
    "        coonv1_2 = self.conv(dec1_2)\n",
    "\n",
    "\n",
    "        \n",
    "        #block3\n",
    "        enc1_3 = self.encoder1(coonv1_2)\n",
    "        enc2_3 = self.encoder2(self.pool1(enc1_3))\n",
    "        enc3_3 = self.encoder3(self.pool2(enc2_3))\n",
    "        enc4_3 = self.encoder4(self.pool3(enc3_3))\n",
    "\n",
    "        bottleneck_3 = self.bottleneck(self.pool4(enc4_3))\n",
    "\n",
    "        dec4_3 = self.upconv4(bottleneck_3)\n",
    "        dec4_3 = torch.cat((dec4_3, enc4_3), dim=1)\n",
    "        dec4_3 = self.decoder4(dec4_3)\n",
    "        dec3_3 = self.upconv3(dec4_3)\n",
    "        dec3_3 = torch.cat((dec3_3, enc3_3), dim=1)\n",
    "        dec3_3 = self.decoder3(dec3_3)\n",
    "        dec2_3 = self.upconv2(dec3_3)\n",
    "        dec2_3 = torch.cat((dec2_3, enc2_3), dim=1)\n",
    "        dec2_3 = self.decoder2(dec2_3)\n",
    "        dec1_3 = self.upconv1(dec2_3)\n",
    "        dec1_3 = torch.cat((dec1_3, enc1_3), dim=1)\n",
    "        dec1_3 = self.decoder1(dec1_3)\n",
    "        coonv1_3 = self.conv(dec1_3)\n",
    "\n",
    "\n",
    "        #block4\n",
    "        enc1_4 = self.encoder1(coonv1_3)\n",
    "        enc2_4 = self.encoder2(self.pool1(enc1_4))\n",
    "        enc3_4 = self.encoder3(self.pool2(enc2_4))\n",
    "        enc4_4 = self.encoder4(self.pool3(enc3_4))\n",
    "\n",
    "        bottleneck_4 = self.bottleneck(self.pool4(enc4_4))\n",
    "\n",
    "        dec4_4 = self.upconv4(bottleneck_4)\n",
    "        dec4_4 = torch.cat((dec4_4, enc4_4), dim=1)\n",
    "        dec4_4 = self.decoder4(dec4_4)\n",
    "        dec3_4 = self.upconv3(dec4_4)\n",
    "        dec3_4 = torch.cat((dec3_4, enc3_4), dim=1)\n",
    "        dec3_4 = self.decoder3(dec3_4)\n",
    "        dec2_4 = self.upconv2(dec3_4)\n",
    "        dec2_4 = torch.cat((dec2_4, enc2_4), dim=1)\n",
    "        dec2_4 = self.decoder2(dec2_4)\n",
    "        dec1_4 = self.upconv1(dec2_4)\n",
    "        dec1_4 = torch.cat((dec1_4, enc1_4), dim=1)\n",
    "        dec1_4 = self.decoder1(dec1_4)\n",
    "        coonv1_4 = self.conv(dec1_4)\n",
    "\n",
    "        #concatenation of different UNet feature maps\n",
    "        concat = torch.cat((coonv1_1, coonv1_2, coonv1_3, coonv1_4  ), dim=1)\n",
    "\n",
    "\n",
    "        #pyramid convolution\n",
    "        conv_pyramid_3 =  self.pyramid_3(concat)\n",
    "        conv_pyramid_5 =  self.pyramid_5(concat)\n",
    "        conv_pyramid_7 =  self.pyramid_7(concat)\n",
    "        conv_pyramid_11 =  self.pyramid_11(concat)\n",
    "        conv_pyramid_17 =  self.pyramid_17(concat)\n",
    "        conv_pyramid_25 =  self.pyramid_25(concat)\n",
    "        conv_pyramid_35 =  self.pyramid_35(concat)\n",
    "        conv_pyramid_45 =  self.pyramid_45(concat)\n",
    "\n",
    "        #concatenation of feature maps corresponding different convolution filters present in Pyramid convolution layer\n",
    "        concat_py = torch.cat(( conv_pyramid_3, conv_pyramid_5, conv_pyramid_7, conv_pyramid_11, conv_pyramid_17, conv_pyramid_25, conv_pyramid_35, conv_pyramid_45), dim=1)\n",
    "      \n",
    "        return torch.sigmoid(self.conv_1(concat_py))\n",
    "\n",
    "    @staticmethod\n",
    "    def _block(in_channels, features, name):\n",
    "        return nn.Sequential(\n",
    "            OrderedDict(\n",
    "                [\n",
    "                    (\n",
    "                        name + \"conv1\",\n",
    "                        nn.Conv2d(\n",
    "                            in_channels=in_channels,\n",
    "                            out_channels=features,\n",
    "                            kernel_size=3,\n",
    "                            padding=1,\n",
    "                            bias=False,\n",
    "                        ),\n",
    "                    ),\n",
    "                    (name + \"norm1\", nn.BatchNorm2d(num_features=features)),\n",
    "                    (name + \"relu1\", nn.ReLU(inplace=True)),\n",
    "                    (\n",
    "                        name + \"conv2\",\n",
    "                        nn.Conv2d(\n",
    "                            in_channels=features,\n",
    "                            out_channels=features,\n",
    "                            kernel_size=3,\n",
    "                            padding=1,\n",
    "                            bias=False,\n",
    "                        ),\n",
    "                    ),\n",
    "                    (name + \"norm2\", nn.BatchNorm2d(num_features=features)),\n",
    "                    (name + \"relu2\", nn.ReLU(inplace=True)),\n",
    "                ]\n",
    "            )\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZzLeMds8bglC"
   },
   "source": [
    "Defining our Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Zk5I8Q_mAAAm"
   },
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, in_channels, use_sigmoid=True, use_spectral_norm=True):\n",
    "        super().__init__()\n",
    "        self.use_sigmoid = use_sigmoid\n",
    "\n",
    "        self.conv1 = self.features = nn.Sequential(\n",
    "            spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "        )\n",
    "\n",
    "        self.conv2 = nn.Sequential(\n",
    "            spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "        )\n",
    "\n",
    "        self.conv3 = nn.Sequential(\n",
    "            spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "        )\n",
    "\n",
    "        self.conv4 = nn.Sequential(\n",
    "            spectral_norm(nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "        )\n",
    "\n",
    "        self.conv5 = nn.Sequential(\n",
    "            spectral_norm(nn.Conv2d(in_channels=512, out_channels=1, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        conv1 = self.conv1(x)\n",
    "        conv2 = self.conv2(conv1)\n",
    "        conv3 = self.conv3(conv2)\n",
    "        conv4 = self.conv4(conv3)\n",
    "        conv5 = self.conv5(conv4)\n",
    "\n",
    "        outputs = conv5\n",
    "        if self.use_sigmoid:\n",
    "            outputs = torch.sigmoid(conv5)\n",
    "\n",
    "        return outputs, [conv1, conv2, conv3, conv4, conv5]\n",
    "\n",
    "def spectral_norm(module, mode=True):\n",
    "    if mode:\n",
    "        return nn.utils.spectral_norm(module)\n",
    "\n",
    "    return module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pqRl8shQbkm_"
   },
   "source": [
    "Defining our loss functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0FLH3Z9lAAAo"
   },
   "outputs": [],
   "source": [
    "class AdversarialLoss(nn.Module):\n",
    "    r\"\"\"\n",
    "    Adversarial loss\n",
    "    https://arxiv.org/abs/1711.10337\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, type='hinge', target_real_label=1.0, target_fake_label=0.0):\n",
    "        r\"\"\"\n",
    "        type = nsgan | lsgan | hinge\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        self.type = type\n",
    "        self.register_buffer('real_label', torch.tensor(target_real_label))\n",
    "        self.register_buffer('fake_label', torch.tensor(target_fake_label))\n",
    "\n",
    "        if type == 'nsgan':\n",
    "            self.criterion = nn.BCELoss()\n",
    "\n",
    "        elif type == 'lsgan':\n",
    "            self.criterion = nn.MSELoss()\n",
    "\n",
    "        elif type == 'hinge':\n",
    "            self.criterion = nn.ReLU()\n",
    "\n",
    "    def __call__(self, outputs, is_real, is_disc=None):\n",
    "        if self.type == 'hinge':\n",
    "            if is_real:\n",
    "              return -torch.log(outputs).mean()\n",
    "            else:\n",
    "              return -torch.log(1-outputs).mean()\n",
    "\n",
    "\n",
    "        else:\n",
    "            labels = (self.real_label if is_real else self.fake_label).expand_as(outputs)\n",
    "            loss = self.criterion(outputs, labels)\n",
    "            return loss\n",
    "\n",
    "\n",
    "\n",
    "class ContentLoss(nn.Module):\n",
    "    r\"\"\"\n",
    "    Perceptual loss, VGG-based\n",
    "    https://arxiv.org/abs/1603.08155\n",
    "    https://github.com/dxyang/StyleTransfer/blob/master/utils.py\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, weights=[1.0, 1.0, 1.0, 1.0, 1.0]):\n",
    "        super().__init__()\n",
    "        self.add_module('vgg', VGG19().cuda())\n",
    "        self.criterion = torch.nn.L1Loss().cuda()\n",
    "        self.weights = weights\n",
    "\n",
    "    def __call__(self, x, y):\n",
    "        # Compute features\n",
    "        x_vgg, y_vgg = self.vgg(x), self.vgg(y)\n",
    "\n",
    "        content_loss = 0.0\n",
    "        content_loss += self.weights[0] * self.criterion(x_vgg['relu1_1'], y_vgg['relu1_1'])\n",
    "        content_loss += self.weights[1] * self.criterion(x_vgg['relu2_1'], y_vgg['relu2_1'])\n",
    "        content_loss += self.weights[2] * self.criterion(x_vgg['relu3_1'], y_vgg['relu3_1'])\n",
    "        content_loss += self.weights[3] * self.criterion(x_vgg['relu4_1'], y_vgg['relu4_1'])\n",
    "        content_loss += self.weights[4] * self.criterion(x_vgg['relu5_1'], y_vgg['relu5_1'])\n",
    "\n",
    "\n",
    "        return content_loss\n",
    "\n",
    "\n",
    "\n",
    "class VGG19(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        features = models.vgg19(pretrained=True).features\n",
    "        self.relu1_1 = torch.nn.Sequential()\n",
    "        self.relu1_2 = torch.nn.Sequential()\n",
    "\n",
    "        self.relu2_1 = torch.nn.Sequential()\n",
    "        self.relu2_2 = torch.nn.Sequential()\n",
    "\n",
    "        self.relu3_1 = torch.nn.Sequential()\n",
    "        self.relu3_2 = torch.nn.Sequential()\n",
    "        self.relu3_3 = torch.nn.Sequential()\n",
    "        self.relu3_4 = torch.nn.Sequential()\n",
    "\n",
    "        self.relu4_1 = torch.nn.Sequential()\n",
    "        self.relu4_2 = torch.nn.Sequential()\n",
    "        self.relu4_3 = torch.nn.Sequential()\n",
    "        self.relu4_4 = torch.nn.Sequential()\n",
    "\n",
    "        self.relu5_1 = torch.nn.Sequential()\n",
    "        self.relu5_2 = torch.nn.Sequential()\n",
    "        self.relu5_3 = torch.nn.Sequential()\n",
    "        self.relu5_4 = torch.nn.Sequential()\n",
    "\n",
    "        for x in range(2):\n",
    "            self.relu1_1.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(2, 4):\n",
    "            self.relu1_2.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(4, 7):\n",
    "            self.relu2_1.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(7, 9):\n",
    "            self.relu2_2.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(9, 12):\n",
    "            self.relu3_1.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(12, 14):\n",
    "            self.relu3_2.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(14, 16):\n",
    "            self.relu3_3.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(16, 18):\n",
    "            self.relu3_4.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(18, 21):\n",
    "            self.relu4_1.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(21, 23):\n",
    "            self.relu4_2.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(23, 25):\n",
    "            self.relu4_3.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(25, 27):\n",
    "            self.relu4_4.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(27, 30):\n",
    "            self.relu5_1.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(30, 32):\n",
    "            self.relu5_2.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(32, 34):\n",
    "            self.relu5_3.add_module(str(x), features[x])\n",
    "\n",
    "        for x in range(34, 36):\n",
    "            self.relu5_4.add_module(str(x), features[x])\n",
    "\n",
    "        # don't need the gradients, just want the features\n",
    "        for param in self.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "    def forward(self, x):\n",
    "        relu1_1 = self.relu1_1(x)\n",
    "        relu1_2 = self.relu1_2(relu1_1)\n",
    "\n",
    "        relu2_1 = self.relu2_1(relu1_2)\n",
    "        relu2_2 = self.relu2_2(relu2_1)\n",
    "\n",
    "        relu3_1 = self.relu3_1(relu2_2)\n",
    "        relu3_2 = self.relu3_2(relu3_1)\n",
    "        relu3_3 = self.relu3_3(relu3_2)\n",
    "        relu3_4 = self.relu3_4(relu3_3)\n",
    "\n",
    "        relu4_1 = self.relu4_1(relu3_4)\n",
    "        relu4_2 = self.relu4_2(relu4_1)\n",
    "        relu4_3 = self.relu4_3(relu4_2)\n",
    "        relu4_4 = self.relu4_4(relu4_3)\n",
    "\n",
    "        relu5_1 = self.relu5_1(relu4_4)\n",
    "        relu5_2 = self.relu5_2(relu5_1)\n",
    "        relu5_3 = self.relu5_3(relu5_2)\n",
    "        relu5_4 = self.relu5_4(relu5_3)\n",
    "\n",
    "        out = {\n",
    "            'relu1_1': relu1_1,\n",
    "            'relu1_2': relu1_2,\n",
    "\n",
    "            'relu2_1': relu2_1,\n",
    "            'relu2_2': relu2_2,\n",
    "\n",
    "            'relu3_1': relu3_1,\n",
    "            'relu3_2': relu3_2,\n",
    "            'relu3_3': relu3_3,\n",
    "            'relu3_4': relu3_4,\n",
    "\n",
    "            'relu4_1': relu4_1,\n",
    "            'relu4_2': relu4_2,\n",
    "            'relu4_3': relu4_3,\n",
    "            'relu4_4': relu4_4,\n",
    "\n",
    "            'relu5_1': relu5_1,\n",
    "            'relu5_2': relu5_2,\n",
    "            'relu5_3': relu5_3,\n",
    "            'relu5_4': relu5_4,\n",
    "        }\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nCVcqVjAAAAr"
   },
   "outputs": [],
   "source": [
    "def gaussian(window_size, sigma):\n",
    "    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])\n",
    "    return gauss/gauss.sum()\n",
    "\n",
    "def create_window(window_size, channel):\n",
    "    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)\n",
    "    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)\n",
    "    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())\n",
    "    return window\n",
    "\n",
    "def _ssim(img1, img2, window, window_size, channel, size_average = True):\n",
    "    mu1 = F3.conv2d(img1, window, padding = window_size//2, groups = channel)\n",
    "    mu2 = F3.conv2d(img2, window, padding = window_size//2, groups = channel)\n",
    "\n",
    "    mu1_sq = mu1.pow(2)\n",
    "    mu2_sq = mu2.pow(2)\n",
    "    mu1_mu2 = mu1*mu2\n",
    "\n",
    "    sigma1_sq = F3.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq\n",
    "    sigma2_sq = F3.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq\n",
    "    sigma12 = F3.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2\n",
    "\n",
    "    C1 = 0.01**2\n",
    "    C2 = 0.03**2\n",
    "\n",
    "    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))\n",
    "\n",
    "    if size_average:\n",
    "        return ssim_map.mean()\n",
    "    else:\n",
    "        return ssim_map.mean(1).mean(1).mean(1)\n",
    "\n",
    "class SSIM(torch.nn.Module):\n",
    "    def __init__(self, window_size = 11, size_average = True):\n",
    "        super(SSIM, self).__init__()\n",
    "        self.window_size = window_size\n",
    "        self.size_average = size_average\n",
    "        self.channel = 1\n",
    "        self.window = create_window(window_size, self.channel)\n",
    "\n",
    "    def forward(self, img1, img2):\n",
    "        (_, channel, _, _) = img1.size()\n",
    "\n",
    "        if channel == self.channel and self.window.data.type() == img1.data.type():\n",
    "            window = self.window\n",
    "        else:\n",
    "            window = create_window(self.window_size, channel)\n",
    "            \n",
    "            if img1.is_cuda:\n",
    "                window = window.cuda(img1.get_device())\n",
    "            window = window.type_as(img1)\n",
    "            \n",
    "            self.window = window\n",
    "            self.channel = channel\n",
    "\n",
    "\n",
    "        return _ssim(img1, img2, window, self.window_size, channel, self.size_average)\n",
    "\n",
    "def ssim(img1, img2, window_size = 11, size_average = True):\n",
    "    (_, channel, _, _) = img1.size()\n",
    "    window = create_window(window_size, channel)\n",
    "    \n",
    "    if img1.is_cuda:\n",
    "        window = window.cuda(img1.get_device())\n",
    "    window = window.type_as(img1)\n",
    "    \n",
    "    return _ssim(img1, img2, window, window_size, channel, size_average)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KtvXXCtxboWU"
   },
   "source": [
    "Defining our model i.e. DU_Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SNTpEUunAAAw"
   },
   "outputs": [],
   "source": [
    "class DU_Net(nn.Module):\n",
    "\n",
    "    def __init__(self, unet_input, unet_output, discriminator_input):\n",
    "        super().__init__()\n",
    "\n",
    "\n",
    "        unet = UNet(in_channels=unet_input ,out_channels=unet_output)\n",
    "        unet = nn.DataParallel(unet, device_ids=[0, 1])\n",
    "        unet = unet.cuda()\n",
    "\n",
    "        discriminator = Discriminator(in_channels=discriminator_input , use_sigmoid=True)\n",
    "        discriminator = nn.DataParallel(discriminator, device_ids=[0, 1])\n",
    "        discriminator = discriminator.cuda()\n",
    "\n",
    "        criterion = nn.MSELoss()\n",
    "        adversarial_loss = AdversarialLoss(type='hinge')\n",
    "        l1_loss = nn.L1Loss()\n",
    "        content_loss = ContentLoss()\n",
    "        ssim = SSIM(window_size = 11)\n",
    "        bce = nn.BCELoss()\n",
    "\n",
    "        self.add_module('unet', unet)\n",
    "        self.add_module('discriminator', discriminator)\n",
    "\n",
    "        self.add_module('criterion', criterion)\n",
    "        self.add_module('adversarial_loss', adversarial_loss)\n",
    "        self.add_module('l1_loss', l1_loss)\n",
    "        self.add_module('content_loss', content_loss)\n",
    "        self.add_module('ssim_loss', ssim)\n",
    "        self.add_module('bce_loss', bce)\n",
    "        \n",
    "\n",
    "        self.unet_optimizer = optim.Adam(\n",
    "            unet.parameters(), \n",
    "            lr = float(0.00001),\n",
    "            betas=(0.9, 0.999)\n",
    "            )\n",
    "\n",
    "        self.dis_optimizer = optim.Adam(\n",
    "             params=discriminator.parameters(),\n",
    "             lr=float(0.00001),\n",
    "             betas=(0.9, 0.999)\n",
    "             )\n",
    "\n",
    "        self.unet_input = unet_input\n",
    "        self.unet_output = unet_output\n",
    "        self.discriminator_input = discriminator_input\n",
    "\n",
    "\n",
    "    def load(self, path_unet, path_discriminator):\n",
    "        weight_unet = torch.load(path_unet)\n",
    "        weight_discriminator = torch.load(path_discriminator)\n",
    "        self.unet.load_state_dict(weight_unet)\n",
    "        self.discriminator.load_state_dict(weight_discriminator)\n",
    "\n",
    "    def save_weight(self, path_unet, path_dis):\n",
    "        torch.save(self.unet.state_dict(), path_unet)\n",
    "        torch.save(self.discriminator.state_dict(), path_dis)\n",
    "\n",
    "    def process(self, haze_images, dehaze_images):\n",
    "\n",
    "        # zero optimizers\n",
    "        self.unet_optimizer.zero_grad()\n",
    "        self.dis_optimizer.zero_grad()\n",
    "\n",
    "\n",
    "        # find output and initialize loss to zero\n",
    "        unet_loss = 0\n",
    "        dis_loss = 0\n",
    "\n",
    "        outputs = self.unet(haze_images.cuda())\n",
    "\n",
    "\n",
    "        # discriminator loss\n",
    "        dis_real, dis_real_feat = self.discriminator(dehaze_images.cuda())        \n",
    "        dis_fake, dis_fake_feat = self.discriminator(outputs.detach().cuda())       \n",
    "        dis_real_loss = self.adversarial_loss(dis_real, True, True)\n",
    "        dis_fake_loss = self.adversarial_loss(dis_fake, False, True)\n",
    "        dis_loss += (dis_real_loss + dis_fake_loss) / 2\n",
    "\n",
    "\n",
    "        # unet loss\n",
    "        unet_fake, unet_fake_feat = self.discriminator(outputs.cuda())        \n",
    "        unet_gan_loss = self.adversarial_loss(unet_fake, True, False) * 0.7\n",
    "        unet_loss += unet_gan_loss\n",
    "\n",
    "        unet_criterion = self.criterion(outputs.cuda(), dehaze_images.cuda())\n",
    "        unet_loss += unet_criterion\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        gen_content_loss = self.content_loss(outputs.cuda(), dehaze_images.cuda())\n",
    "        gen_content_loss = (gen_content_loss * 0.7).cuda()\n",
    "        unet_loss += gen_content_loss.cuda()\n",
    "        \n",
    "        \n",
    "        ssim_loss =  self.ssim_loss(outputs.cuda(), dehaze_images.cuda())\n",
    "        ssim_loss = (1-ssim_loss)*2\n",
    "        unet_loss += ssim_loss.cuda()\n",
    "\n",
    "        return unet_loss, dis_loss, unet_criterion, 1-ssim_loss/2\n",
    "\n",
    "    def backward(self, unet_loss, dis_loss):\n",
    "        dis_loss.backward(retain_graph = True)\n",
    "        self.dis_optimizer.step()\n",
    "\n",
    "        unet_loss.backward()\n",
    "        self.unet_optimizer.step()\n",
    "        \n",
    "\n",
    "    def predict(self, haze_images):\n",
    "      predict_mask = self.unet(haze_images.cuda())\n",
    "      return predict_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "d5NG1aLzbtmE"
   },
   "source": [
    "Defining our and creating the data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "B4Zt47MJAAA0"
   },
   "outputs": [],
   "source": [
    "class Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, haze_list, dehaze_list, augment=False):\n",
    "        super().__init__()\n",
    "        self.augment = augment\n",
    "        self.haze_list = haze_list\n",
    "        self.dehaze_list = dehaze_list\n",
    "        \n",
    "    def __len__(self):\n",
    "        return 210\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        try:\n",
    "            item = self.load_item(index)\n",
    "        except:\n",
    "            print('loading error: ' + self.haze_list[index])\n",
    "            item = self.load_item(0)\n",
    "            \n",
    "\n",
    "        return item\n",
    "\n",
    "\n",
    "    def load_item(self, index):\n",
    "        val = 1024*2           #crop size i.e hight and width\n",
    "        size_data = 25         #depends on the no. of training images in the dataset\n",
    "        height_data = 4657     #heigth of the training images\n",
    "        width_data = 2833      #width of the training images\n",
    "\n",
    "        numx = random.randint(0, height_data-val)\n",
    "        numy = random.randint(0, width_data-val)\n",
    "\n",
    "        haze_image = cv2.imread(self.haze_list[index%size_data])\n",
    "        dehaze_image = cv2.imread(self.dehaze_list[index%size_data])\n",
    "        haze_image = Image.fromarray(haze_image)\n",
    "        dehaze_image = Image.fromarray(dehaze_image)\n",
    "\n",
    "        haze_crop=haze_image.crop((numx, numy, numx+val, numy+val))\n",
    "        dehaze_crop=dehaze_image.crop((numx, numy, numx+val, numy+val))\n",
    " \n",
    "        haze_crop = haze_crop.resize((512,512), resample=PIL.Image.BICUBIC)\n",
    "        dehaze_crop = dehaze_crop.resize((512,512), resample=PIL.Image.BICUBIC)\n",
    "\n",
    "        haze_crop = np.array(haze_crop)\n",
    "        dehaze_crop = np.array(dehaze_crop)\n",
    "        haze_crop = cv2.cvtColor(haze_crop, cv2.COLOR_BGR2YCrCb)\n",
    "        dehaze_crop = cv2.cvtColor(dehaze_crop, cv2.COLOR_BGR2YCrCb)\n",
    "        haze_crop = self.to_tensor(haze_crop).cuda()\n",
    "        dehaze_crop = self.to_tensor(dehaze_crop).cuda()\n",
    "        \n",
    "        return haze_crop.cuda(), dehaze_crop.cuda()\n",
    "    \n",
    "    def to_tensor(self, img):\n",
    "        img_t = F.to_tensor(img).float()\n",
    "        return img_t\n",
    "\n",
    "\n",
    "    def create_iterator(self, batch_size):\n",
    "        while True:\n",
    "            sample_loader = DataLoader(\n",
    "                dataset=self,\n",
    "                batch_size=batch_size,\n",
    "                drop_last=True\n",
    "            )\n",
    "\n",
    "            for item in sample_loader:\n",
    "                yield item\n",
    "\n",
    "\n",
    "path_of_train_hazy_images = 'train/haze/*.png'\n",
    "path_of_train_gt_images = 'train/gt/*.png'\n",
    "\n",
    "images_paths_train_gt=glob.glob(path_of_train_gt_images)\n",
    "image_paths_train_hazy=glob.glob(path_of_train_hazy_images)\n",
    "\n",
    "train_dataset = Dataset(image_paths_train_hazy, images_paths_train_gt, augment=False)\n",
    "\n",
    "train_loader = DataLoader(\n",
    "            dataset=train_dataset,\n",
    "            batch_size=2,\n",
    "            num_workers=0,\n",
    "            drop_last=True,\n",
    "            shuffle=False\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DQW0X3g5dNBX"
   },
   "source": [
    "Creating the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZACwe8ObAABE"
   },
   "outputs": [],
   "source": [
    "graph_gloss = []\n",
    "input_unet_channel = 3\n",
    "output_unet_channel = 3\n",
    "input_dis_channel = 3\n",
    "max_epochs = 100\n",
    "DUNet = DU_Net(input_unet_channel ,output_unet_channel ,input_dis_channel).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Gb7PvZASdQn1"
   },
   "source": [
    "Training function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "24-9evbnAABM"
   },
   "outputs": [],
   "source": [
    "def train(max_epochs):\n",
    "    for epoch in range(max_epochs):\n",
    "        i=1\n",
    "        mse_epoch = 0.0\n",
    "        ssim_epoch = 0.0\n",
    "        unet_epoch = 0.0\n",
    "        for haze_images, dehaze_images, in train_loader:\n",
    "            unet_loss, dis_loss, mse, ssim = DUNet.process(haze_images.cuda(), dehaze_images.cuda())\n",
    "            DUNet.backward(unet_loss.cuda(), dis_loss.cuda())\n",
    "            print('Epoch: '+str(epoch+1)+ ' || Batch: '+str(i)+ \" || unet loss: \"+str(unet_loss.cpu().item()) + \" || dis loss: \"+str(dis_loss.cpu().item()) + \" || mse: \"+str(mse.cpu().item()) + \" | ssim:\" + str(ssim.cpu().item()) )\n",
    "            mse_epoch =  mse_epoch + mse.cpu().item() \n",
    "            ssim_epoch = ssim_epoch + ssim.cpu().item()\n",
    "            unet_epoch = unet_epoch + unet_loss.cpu().item()\n",
    "            i=i+1\n",
    "        \n",
    "        print()\n",
    "        mse_epoch = mse_epoch/i\n",
    "        ssim_epoch = ssim_epoch/i\n",
    "        unet_epoch = unet_epoch/i\n",
    "        graph_gloss.append(ssim_epoch)\n",
    "        print(\"mse: + \"+str(mse_epoch) + \" | ssim: \"+ str(ssim_epoch)+ \" | unet:\"+str(unet_epoch))\n",
    "        print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "29CCGgMDdTPe"
   },
   "source": [
    "Calling training function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rmEpSm9jAABP"
   },
   "outputs": [],
   "source": [
    "epochs = 150\n",
    "train(epochs) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3vXFZE36R-C_"
   },
   "source": [
    "Saving weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZKT5tEJmAABT"
   },
   "outputs": [],
   "source": [
    "path_of_generator_weight = 'weight/generator.pth'  #path for storing the weights of genertaor\n",
    "path_of_discriminator_weight = 'weight/discriminator.pth'  #path for storing the weights of discriminator\n",
    "DUNet.save_weight(path_of_generator_weight,path_of_discriminator_weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AWqkb8epSIup"
   },
   "source": [
    "Saving weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aPjwtGKeAABV"
   },
   "outputs": [],
   "source": [
    "path_of_generator_weight = 'weight/generator.pth'  #path where the weights of genertaor are stored\n",
    "path_of_discriminator_weight = 'weight/discriminator.pth'  #path where the weights of discriminator are stored\n",
    "DUNet.load('weights/new/in all/u_' + str(21) + '.pth','weights/new/in all/d_' + str(21) + '.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3MAztif9SRC_"
   },
   "source": [
    "Runing the model on test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rEMIiCj3AABe"
   },
   "outputs": [],
   "source": [
    "def to_tensor(img):\n",
    "    img_t = F6.to_tensor(img).float()\n",
    "    return img_t\n",
    "\n",
    "def postprocess(img):\n",
    "        img = img * 255.0\n",
    "        img = img.permute(0, 2, 3, 1)\n",
    "        return img.int()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7Bb-OyJTAABh"
   },
   "outputs": [],
   "source": [
    "path_of_test_hazy_images = 'test/haze/*.png'\n",
    "path_for_resultant_dehaze_images = 'test/result/'\n",
    "image_paths_test_hazy=glob.glob(path_of_test_hazy_images)\n",
    "\n",
    "for i in range(len(image_paths_test_hazy)):\n",
    "    haze_image = cv2.imread(image_paths_test_hazy[i])\n",
    "    haze_image = Image.fromarray(haze_image)\n",
    "    haze_image = haze_image.resize((512,512), resample=PIL.Image.BICUBIC)\n",
    "    haze_image = np.array(haze_image)\n",
    "    haze_image = cv2.cvtColor(haze_image, cv2.COLOR_BGR2YCrCb)\n",
    "    haze_image = to_tensor(haze_image).cuda()\n",
    "    haze_image = haze_image.reshape(1,3,512,512)\n",
    "\n",
    "    dehaze_image = DUNet.predict(haze_image) \n",
    "    \n",
    "    dehaze_image = postprocess(dehaze_image)[0]\n",
    "    dehaze_image = dehaze_image.cpu().detach().numpy()\n",
    "    dehaze_image = dehaze_image.astype('uint8')\n",
    "    dehaze_image = dehaze_image.reshape(512,512,3)\n",
    "    dehaze_image = cv2.cvtColor(dehaze_image, cv2.COLOR_YCrCb2BGR)\n",
    "    cv2.imwrite(path_for_resultant_dehaze_images+str(i+50)+'.png', dehaze_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YcYTCrl_SZGz"
   },
   "source": [
    "Calculating the metrices i.e PSNR and SSIM for testing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jgJC-JMuAABl"
   },
   "outputs": [],
   "source": [
    "class EdgeAccuracy(nn.Module):\n",
    "    \"\"\"\n",
    "    Measures the accuracy of the edge map\n",
    "    \"\"\"\n",
    "    def __init__(self, threshold=0.5):\n",
    "        super().__init__()\n",
    "        self.threshold = threshold\n",
    "\n",
    "    def __call__(self, inputs, outputs):\n",
    "        labels = (inputs > self.threshold)\n",
    "        outputs = (outputs > self.threshold)\n",
    "\n",
    "        relevant = torch.sum(labels.float())\n",
    "        selected = torch.sum(outputs.float())\n",
    "\n",
    "        if relevant == 0 and selected == 0:\n",
    "            return 1, 1\n",
    "\n",
    "        true_positive = ((outputs == labels) * labels).float()\n",
    "        recall = torch.sum(true_positive) / (relevant + 1e-8)\n",
    "        precision = torch.sum(true_positive) / (selected + 1e-8)\n",
    "\n",
    "        return precision, recall\n",
    "\n",
    "\n",
    "class PSNR(nn.Module):\n",
    "    def __init__(self, max_val=0):\n",
    "        super().__init__()\n",
    "\n",
    "        base10 = torch.log(torch.tensor(10.0))\n",
    "        max_val = torch.tensor(max_val).float()\n",
    "\n",
    "        self.register_buffer('base10', base10)\n",
    "        self.register_buffer('max_val', 20 * torch.log(max_val) / base10)\n",
    "\n",
    "    def __call__(self, a, b):\n",
    "        mse = torch.mean((a.float() - b.float()) ** 2)\n",
    "    \n",
    "        if mse == 0:\n",
    "            return 0\n",
    "\n",
    "        return 1.0 / mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Fsxd9gSxAABm"
   },
   "outputs": [],
   "source": [
    "def gaussian(window_size, sigma):\n",
    "    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])\n",
    "    return gauss/gauss.sum()\n",
    "\n",
    "def create_window(window_size, channel):\n",
    "    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)\n",
    "    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)\n",
    "    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())\n",
    "    return window\n",
    "\n",
    "def _ssim(img1, img2, window, window_size, channel, size_average = True):\n",
    "    mu1 = F9.conv2d(img1, window, padding = window_size//2, groups = channel)\n",
    "    mu2 = F9.conv2d(img2, window, padding = window_size//2, groups = channel)\n",
    "\n",
    "    mu1_sq = mu1.pow(2)\n",
    "    mu2_sq = mu2.pow(2)\n",
    "    mu1_mu2 = mu1*mu2\n",
    "\n",
    "    sigma1_sq = F9.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq\n",
    "    sigma2_sq = F9.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq\n",
    "    sigma12 = F9.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2\n",
    "\n",
    "    C1 = 0.01**2\n",
    "    C2 = 0.03**2\n",
    "\n",
    "    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))\n",
    "\n",
    "    if size_average:\n",
    "        return ssim_map.mean()\n",
    "    else:\n",
    "        return ssim_map.mean(1).mean(1).mean(1)\n",
    "\n",
    "class SSIM(torch.nn.Module):\n",
    "    def __init__(self, window_size = 11, size_average = True):\n",
    "        super(SSIM, self).__init__()\n",
    "        self.window_size = window_size\n",
    "        self.size_average = size_average\n",
    "        self.channel = 1\n",
    "        self.window = create_window(window_size, self.channel)\n",
    "\n",
    "    def forward(self, img1, img2):\n",
    "        (_, channel, _, _) = img1.size()\n",
    "\n",
    "        if channel == self.channel and self.window.data.type() == img1.data.type():\n",
    "            window = self.window\n",
    "        else:\n",
    "            window = create_window(self.window_size, channel)\n",
    "            \n",
    "            if img1.is_cuda:\n",
    "                window = window.cuda(img1.get_device())\n",
    "            window = window.type_as(img1)\n",
    "            \n",
    "            self.window = window\n",
    "            self.channel = channel\n",
    "\n",
    "\n",
    "        return _ssim(img1, img2, window, self.window_size, channel, self.size_average)\n",
    "\n",
    "def ssim(img1, img2, window_size = 11, size_average = True):\n",
    "    (_, channel, _, _) = img1.size()\n",
    "    window = create_window(window_size, channel)\n",
    "    \n",
    "    if img1.is_cuda:\n",
    "        window = window.cuda(img1.get_device())\n",
    "    window = window.type_as(img1)\n",
    "    \n",
    "    return _ssim(img1, img2, window, window_size, channel, size_average)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9hTzS1wUAABq"
   },
   "outputs": [],
   "source": [
    "ssim = SSIM(window_size = 11)\n",
    "psnr = PSNR()\n",
    "psnr_val = 0\n",
    "psnr_val = 0.0\n",
    "final_ssim = 0\n",
    "\n",
    "path_of_test_hazy_images = 'test/haze/*.png'\n",
    "path_of_test_gt_images = 'test/gt/*.png'\n",
    "path_for_resultant_dehaze_images = 'test/result/'\n",
    "\n",
    "image_paths_test_hazy=glob.glob(path_of_test_hazy_images)\n",
    "image_paths_test_gt=glob.glob(path_of_test_gt_images)\n",
    "\n",
    "for i in range(len(image_paths_test_hazy)):\n",
    "    im1 = cv2.imread(image_paths_GT[i])\n",
    "    im1 = Image.fromarray(im1)\n",
    "    im1 = im1.resize((512,512), resample=PIL.Image.BICUBIC)\n",
    "    im1 = np.array(im1)\n",
    "    im2 = cv2.imread('Results/experiment yrcrcb/new/outdoor/' + str(i+50)+'.png')\n",
    "\n",
    "    im1 = to_tensor(im1).reshape(1,3,512,512)\n",
    "    im2 = to_tensor(im2).reshape(1,3,512,512)\n",
    "    \n",
    "    psnr_val = psnr(im1, im2)\n",
    "    final_psnr = final_psnr + 10*np.log10((psnr_val))\n",
    "    final_ssim = final_ssim + ssim(im1, im2)\n",
    "\n",
    "\n",
    "print(final_ssim/5.0, final_psnr/5.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WGWLUya4AAB0"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Single Image Dehazing.ipynb",
   "provenance": []
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}