3.Methods/pix2pix_architecture.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "pix2pix_architecture.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "k79o1BFFRYW1"
      },
      "source": [
        "from __future__ import print_function\r\n",
        "import numpy as np\r\n",
        "import math\r\n",
        "import scipy\r\n",
        "import pandas as pd\r\n",
        "import PIL\r\n",
        "import gdal\r\n",
        "import matplotlib.pyplot as plt\r\n",
        "plt.style.use('ggplot')\r\n",
        "import sys, os\r\n",
        "from pathlib import Path\r\n",
        "import time\r\n",
        "import xml.etree.ElementTree as ET\r\n",
        "import random\r\n",
        "import collections, functools, operator\r\n",
        "import csv\r\n",
        "\r\n",
        "import ee\r\n",
        "\r\n",
        "from osgeo import gdal,osr\r\n",
        "from gdalconst import *\r\n",
        "import subprocess\r\n",
        "from osgeo.gdalconst import GA_Update\r\n",
        "\r\n",
        "import torch\r\n",
        "import torch.nn as nn\r\n",
        "from torch.utils.data import Dataset\r\n",
        "from torch.autograd import Variable\r\n",
        "from torch.nn import Linear, ReLU, CrossEntropyLoss, MSELoss, Sequential, Conv2d, MaxPool2d, Module, Softmax, BatchNorm2d, Dropout, Sigmoid\r\n",
        "from torch.optim import Adam, SGD\r\n",
        "from torchvision import transforms, utils\r\n",
        "\r\n",
        "import skimage\r\n",
        "from skimage import io, transform\r\n",
        "import sklearn\r\n",
        "import sklearn.metrics\r\n",
        "from sklearn.feature_extraction import image\r\n",
        "from sklearn import svm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cC7FYGdRRj_o"
      },
      "source": [
        "# P2P architecture"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M3R_yWgXRZEk"
      },
      "source": [
        "'''\r\n",
        "P2P architecture code is based on deeplearning.ai's architecture as defined in the GANs specialization \r\n",
        "'''\r\n",
        "class ContractingBlock(nn.Module):\r\n",
        "    '''\r\n",
        "    ContractingBlock Class\r\n",
        "    Performs two convolutions followed by a max pool operation.\r\n",
        "    Values:\r\n",
        "        input_channels: the number of channels to expect from a given input\r\n",
        "    '''\r\n",
        "    def __init__(self, input_channels, use_dropout=False, use_bn=True):\r\n",
        "        super(ContractingBlock, self).__init__()\r\n",
        "        self.conv1 = nn.Conv2d(input_channels, input_channels * 2, kernel_size=3, padding=1)\r\n",
        "        self.conv2 = nn.Conv2d(input_channels * 2, input_channels * 2, kernel_size=3, padding=1)\r\n",
        "        self.activation = nn.LeakyReLU(0.2)\r\n",
        "        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)\r\n",
        "        if use_bn:\r\n",
        "            self.batchnorm = nn.BatchNorm2d(input_channels * 2)\r\n",
        "        self.use_bn = use_bn\r\n",
        "        if use_dropout:\r\n",
        "            self.dropout = nn.Dropout()\r\n",
        "        self.use_dropout = use_dropout\r\n",
        "\r\n",
        "    def forward(self, x):\r\n",
        "        '''\r\n",
        "        Function for completing a forward pass of ContractingBlock: \r\n",
        "        Given an image tensor, completes a contracting block and returns the transformed tensor.\r\n",
        "        Parameters:\r\n",
        "            x: image tensor of shape (batch size, channels, height, width)\r\n",
        "        '''\r\n",
        "        x = self.conv1(x)\r\n",
        "        if self.use_bn:\r\n",
        "            x = self.batchnorm(x)\r\n",
        "        if self.use_dropout:\r\n",
        "            x = self.dropout(x)\r\n",
        "        x = self.activation(x)\r\n",
        "        x = self.conv2(x)\r\n",
        "        if self.use_bn:\r\n",
        "            x = self.batchnorm(x)\r\n",
        "        if self.use_dropout:\r\n",
        "            x = self.dropout(x)\r\n",
        "        x = self.activation(x)\r\n",
        "        x = self.maxpool(x)\r\n",
        "        return x\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "class ExpandingBlock(nn.Module):\r\n",
        "    '''\r\n",
        "    ExpandingBlock Class:\r\n",
        "    Performs an upsampling, a convolution, a concatenation of its two inputs,\r\n",
        "    followed by two more convolutions with optional dropout\r\n",
        "    Values:\r\n",
        "        input_channels: the number of channels to expect from a given input\r\n",
        "    '''\r\n",
        "    def __init__(self, input_channels, use_dropout=False, use_bn=True):\r\n",
        "        super(ExpandingBlock, self).__init__()\r\n",
        "        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\r\n",
        "        self.conv1 = nn.Conv2d(input_channels, input_channels // 2, kernel_size=2)\r\n",
        "        self.conv2 = nn.Conv2d(input_channels, input_channels // 2, kernel_size=3, padding=1)\r\n",
        "        self.conv3 = nn.Conv2d(input_channels // 2, input_channels // 2, kernel_size=2, padding=1)\r\n",
        "        if use_bn:\r\n",
        "            self.batchnorm = nn.BatchNorm2d(input_channels // 2)\r\n",
        "        self.use_bn = use_bn\r\n",
        "        self.activation = nn.ReLU()\r\n",
        "        if use_dropout:\r\n",
        "            self.dropout = nn.Dropout()\r\n",
        "        self.use_dropout = use_dropout\r\n",
        "\r\n",
        "    def forward(self, x, skip_con_x):\r\n",
        "        '''\r\n",
        "        Function for completing a forward pass of ExpandingBlock: \r\n",
        "        Given an image tensor, completes an expanding block and returns the transformed tensor.\r\n",
        "        Parameters:\r\n",
        "            x: image tensor of shape (batch size, channels, height, width)\r\n",
        "            skip_con_x: the image tensor from the contracting path (from the opposing block of x)\r\n",
        "                    for the skip connection\r\n",
        "        '''\r\n",
        "        x = self.upsample(x)\r\n",
        "        x = self.conv1(x)\r\n",
        "        skip_con_x = crop(skip_con_x, x.shape)\r\n",
        "        x = torch.cat([x, skip_con_x], axis=1)\r\n",
        "        x = self.conv2(x)\r\n",
        "        if self.use_bn:\r\n",
        "            x = self.batchnorm(x)\r\n",
        "        if self.use_dropout:\r\n",
        "            x = self.dropout(x)\r\n",
        "        x = self.activation(x)\r\n",
        "        x = self.conv3(x)\r\n",
        "        if self.use_bn:\r\n",
        "            x = self.batchnorm(x)\r\n",
        "        if self.use_dropout:\r\n",
        "            x = self.dropout(x)\r\n",
        "        x = self.activation(x)\r\n",
        "        return x\r\n",
        "\r\n",
        "\r\n",
        "class FeatureMapBlock(nn.Module):\r\n",
        "    '''\r\n",
        "    FeatureMapBlock Class\r\n",
        "    The final layer of a U-Net - \r\n",
        "    maps each pixel to a pixel with the correct number of output dimensions\r\n",
        "    using a 1x1 convolution.\r\n",
        "    Values:\r\n",
        "        input_channels: the number of channels to expect from a given input\r\n",
        "        output_channels: the number of channels to expect for a given output\r\n",
        "    '''\r\n",
        "    def __init__(self, input_channels, output_channels):\r\n",
        "        super(FeatureMapBlock, self).__init__()\r\n",
        "        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=1)\r\n",
        "\r\n",
        "    def forward(self, x):\r\n",
        "        '''\r\n",
        "        Function for completing a forward pass of FeatureMapBlock: \r\n",
        "        Given an image tensor, returns it mapped to the desired number of channels.\r\n",
        "        Parameters:\r\n",
        "            x: image tensor of shape (batch size, channels, height, width)\r\n",
        "        '''\r\n",
        "        x = self.conv(x)\r\n",
        "        return x\r\n",
        "\r\n",
        "\r\n",
        "class UNet(nn.Module):\r\n",
        "    '''\r\n",
        "    UNet Class\r\n",
        "    A series of 4 contracting blocks followed by 4 expanding blocks to \r\n",
        "    transform an input image into the corresponding paired image, with an upfeature\r\n",
        "    layer at the start and a downfeature layer at the end.\r\n",
        "    Values:\r\n",
        "        input_channels: the number of channels to expect from a given input\r\n",
        "        output_channels: the number of channels to expect for a given output\r\n",
        "    '''\r\n",
        "    def __init__(self, input_channels, output_channels, hidden_channels=32):\r\n",
        "        super(UNet, self).__init__()\r\n",
        "        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)\r\n",
        "        self.contract1 = ContractingBlock(hidden_channels, use_dropout=True)\r\n",
        "        self.contract2 = ContractingBlock(hidden_channels * 2, use_dropout=True)\r\n",
        "        self.contract3 = ContractingBlock(hidden_channels * 4, use_dropout=True)\r\n",
        "        self.contract4 = ContractingBlock(hidden_channels * 8)\r\n",
        "        self.contract5 = ContractingBlock(hidden_channels * 16)\r\n",
        "        self.contract6 = ContractingBlock(hidden_channels * 32)\r\n",
        "        self.expand0 = ExpandingBlock(hidden_channels * 64)\r\n",
        "        self.expand1 = ExpandingBlock(hidden_channels * 32)\r\n",
        "        self.expand2 = ExpandingBlock(hidden_channels * 16)\r\n",
        "        self.expand3 = ExpandingBlock(hidden_channels * 8)\r\n",
        "        self.expand4 = ExpandingBlock(hidden_channels * 4)\r\n",
        "        self.expand5 = ExpandingBlock(hidden_channels * 2)\r\n",
        "        self.downfeature = FeatureMapBlock(hidden_channels, output_channels)\r\n",
        "        self.sigmoid = torch.nn.Sigmoid()\r\n",
        "\r\n",
        "    def forward(self, x):\r\n",
        "        '''\r\n",
        "        Function for completing a forward pass of UNet: \r\n",
        "        Given an image tensor, passes it through U-Net and returns the output.\r\n",
        "        Parameters:\r\n",
        "            x: image tensor of shape (batch size, channels, height, width)\r\n",
        "        '''\r\n",
        "        x0 = self.upfeature(x)\r\n",
        "        x1 = self.contract1(x0)\r\n",
        "        x2 = self.contract2(x1)\r\n",
        "        x3 = self.contract3(x2)\r\n",
        "        x4 = self.contract4(x3)\r\n",
        "        x5 = self.contract5(x4)\r\n",
        "        x6 = self.contract6(x5)\r\n",
        "        x7 = self.expand0(x6, x5)\r\n",
        "        x8 = self.expand1(x7, x4)\r\n",
        "        x9 = self.expand2(x8, x3)\r\n",
        "        x10 = self.expand3(x9, x2)\r\n",
        "        x11 = self.expand4(x10, x1)\r\n",
        "        x12 = self.expand5(x11, x0)\r\n",
        "        xn = self.downfeature(x12)\r\n",
        "        return self.sigmoid(xn)\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "class Discriminator(nn.Module):\r\n",
        "    '''\r\n",
        "    Discriminator Class\r\n",
        "    Structured like the contracting path of the U-Net, the discriminator will\r\n",
        "    output a matrix of values classifying corresponding portions of the image as real or fake. \r\n",
        "    Parameters:\r\n",
        "        input_channels: the number of image input channels\r\n",
        "        hidden_channels: the initial number of discriminator convolutional filters\r\n",
        "    '''\r\n",
        "    def __init__(self, input_channels, hidden_channels=8):\r\n",
        "        super(Discriminator, self).__init__()\r\n",
        "        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)\r\n",
        "        self.contract1 = ContractingBlock(hidden_channels, use_bn=False)\r\n",
        "        self.contract2 = ContractingBlock(hidden_channels * 2)\r\n",
        "        self.contract3 = ContractingBlock(hidden_channels * 4)\r\n",
        "        self.contract4 = ContractingBlock(hidden_channels * 8)\r\n",
        "        self.final = nn.Conv2d(hidden_channels * 16, 1, kernel_size=1)\r\n",
        "\r\n",
        "    def forward(self, x, y):\r\n",
        "        x = torch.cat([x, y], axis=1)\r\n",
        "        x0 = self.upfeature(x)\r\n",
        "        x1 = self.contract1(x0)\r\n",
        "        x2 = self.contract2(x1)\r\n",
        "        x3 = self.contract3(x2)\r\n",
        "        x4 = self.contract4(x3)\r\n",
        "        xn = self.final(x4)\r\n",
        "        return xn\r\n",
        "\r\n",
        "\r\n",
        "def crop(image, new_shape):\r\n",
        "    '''\r\n",
        "    Function for cropping an image tensor: Given an image tensor and the new shape,\r\n",
        "    crops to the center pixels.\r\n",
        "    Parameters:\r\n",
        "        image: image tensor of shape (batch size, channels, height, width)\r\n",
        "        new_shape: a torch.Size object with the shape you want x to have\r\n",
        "    '''\r\n",
        "    middle_height = image.shape[2] // 2\r\n",
        "    middle_width = image.shape[3] // 2\r\n",
        "    starting_height = middle_height - round(new_shape[2] / 2)\r\n",
        "    final_height = starting_height + new_shape[2]\r\n",
        "    starting_width = middle_width - round(new_shape[3] / 2)\r\n",
        "    final_width = starting_width + new_shape[3]\r\n",
        "    cropped_image = image[:, :, starting_height:final_height, starting_width:final_width]\r\n",
        "    return cropped_image\r\n",
        "\r\n",
        "def get_gen_loss(gen, disc, real, condition, adv_criterion, recon_criterion, lambda_recon):\r\n",
        "    '''\r\n",
        "    Return the loss of the generator given inputs.\r\n",
        "    Parameters:\r\n",
        "        gen: the generator; takes the condition and returns potential images\r\n",
        "        disc: the discriminator; takes images and the condition and\r\n",
        "          returns real/fake prediction matrices\r\n",
        "        real: the real images (e.g. maps) to be used to evaluate the reconstruction\r\n",
        "        condition: the source images (e.g. satellite imagery) which are used to produce the real images\r\n",
        "        adv_criterion: the adversarial loss function; takes the discriminator \r\n",
        "                  predictions and the true labels and returns a adversarial \r\n",
        "                  loss (which you aim to minimize)\r\n",
        "        recon_criterion: the reconstruction loss function; takes the generator \r\n",
        "                    outputs and the real images and returns a reconstructuion \r\n",
        "                    loss (which you aim to minimize)\r\n",
        "        lambda_recon: the degree to which the reconstruction loss should be weighted in the sum\r\n",
        "    '''\r\n",
        "    fake = gen(condition)\r\n",
        "    disc_fake_hat = disc(fake, condition)\r\n",
        "    gen_adv_loss = adv_criterion(disc_fake_hat, torch.ones_like(disc_fake_hat))\r\n",
        "    gen_rec_loss = recon_criterion(real, fake)\r\n",
        "    gen_loss = gen_adv_loss + lambda_recon * gen_rec_loss\r\n",
        "    \r\n",
        "    return gen_loss\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "def P2Ptrain(save_model=False):\r\n",
        "    mean_generator_loss_list = []\r\n",
        "    mean_discriminator_loss_list = []\r\n",
        "\r\n",
        "    for epoch in range(num_epochs):\r\n",
        "        # Dataloader returns the batches\r\n",
        "        mean_generator_loss = 0\r\n",
        "        mean_discriminator_loss = 0\r\n",
        "        if epoch == 70: # lr: 0.005 => 0.001\r\n",
        "            gen_opt.param_groups[0]['lr'] = 0.001\r\n",
        "            disc_opt.param_groups[0]['lr'] = 0.001\r\n",
        "        \r\n",
        "        for sample in train_loader:\r\n",
        "            condition = sample['input'] # ALI\r\n",
        "            real = sample['target'] # hyperion\r\n",
        "\r\n",
        "            if readFromPatches:\r\n",
        "                condition = condition[0]\r\n",
        "                real = real[0]\r\n",
        "\r\n",
        "            condition = condition.to(device)\r\n",
        "            real = real.to(device)\r\n",
        "\r\n",
        "            ### Update discriminator ###\r\n",
        "            disc_opt.zero_grad() # Zero out the gradient before backpropagation\r\n",
        "            with torch.no_grad():\r\n",
        "                fake = gen(condition)\r\n",
        "            disc_fake_hat = disc(fake.detach(), condition) # Detach generator\r\n",
        "            disc_fake_loss = adv_criterion(disc_fake_hat, torch.zeros_like(disc_fake_hat))\r\n",
        "            disc_real_hat = disc(real, condition)\r\n",
        "            disc_real_loss = adv_criterion(disc_real_hat, torch.ones_like(disc_real_hat))\r\n",
        "            disc_loss = (disc_fake_loss + disc_real_loss) / 2\r\n",
        "            disc_loss.backward(retain_graph=True) # Update gradients\r\n",
        "            disc_opt.step() # Update optimizer\r\n",
        "\r\n",
        "            ### Update generator ###\r\n",
        "            gen_opt.zero_grad()\r\n",
        "            gen_loss = get_gen_loss(gen, disc, real, condition, adv_criterion, recon_criterion, lambda_recon)\r\n",
        "            gen_loss.backward() # Update gradients\r\n",
        "            gen_opt.step() # Update optimizer\r\n",
        "\r\n",
        "            # Keep track of the average discriminator loss\r\n",
        "            mean_discriminator_loss += disc_loss.item()\r\n",
        "            # Keep track of the average generator loss\r\n",
        "            mean_generator_loss += gen_loss.item()\r\n",
        "            \r\n",
        "\r\n",
        "        mean_generator_loss = mean_generator_loss / len(train_loader)\r\n",
        "        mean_discriminator_loss = mean_discriminator_loss / len(train_loader)\r\n",
        "        mean_generator_loss_list.append(mean_generator_loss)\r\n",
        "        mean_discriminator_loss_list.append(mean_discriminator_loss)\r\n",
        "        ### Visualization code ###\r\n",
        "        if epoch % display_epoch == 0:\r\n",
        "            fig, axs = plt.subplots(2,1)\r\n",
        "            axs[0].plot(mean_generator_loss_list)\r\n",
        "            axs[0].set_title('Generator loss')\r\n",
        "            axs[1].plot(mean_discriminator_loss_list)\r\n",
        "            axs[1].set_title('Discriminator loss')\r\n",
        "            plt.show()\r\n",
        "            print(f\"Epoch {epoch}: Generator (U-Net) loss: {mean_generator_loss}, Discriminator loss: {mean_discriminator_loss}\")\r\n",
        "\r\n",
        "            show_patches(condition.cpu(), fake.cpu(), real.cpu())\r\n",
        "            #calc_metrics(real.cpu().numpy(), fake.cpu().numpy())\r\n",
        "\r\n",
        "            if epoch % 20 == 0:\r\n",
        "                torch.save({'gen': gen.state_dict(),\r\n",
        "                            'gen_opt': gen_opt.state_dict(),\r\n",
        "                            'disc': disc.state_dict(),\r\n",
        "                            'disc_opt': disc_opt.state_dict()\r\n",
        "                        }, os.getcwd() + f\"/drive/My Drive/TFG/Models/NewModel/epoch{epoch}.pth\")\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "def P2Ptest(inferenceDataset, vizImages=False, svc=None, saveMetrics=None):\r\n",
        "    metrics = {'PCC': np.array([0.]*170),\r\n",
        "               'RMSE': np.array([0.]*170),\r\n",
        "               'PSNR': np.array([0.]*170),\r\n",
        "               'SSIM': np.array([0.]*170),\r\n",
        "               'SAM': np.array([0.]*64*64),\r\n",
        "               'SID': np.array([0.]*64*64)}\r\n",
        "    for i, sample in enumerate(inferenceDataset):\r\n",
        "        input = sample['input'][0]\r\n",
        "        prediction = gen(input.to(device)).detach().cpu().numpy()\r\n",
        "        target = sample['target'][0].numpy()\r\n",
        "\r\n",
        "        #VISUALIZATION\r\n",
        "        if vizImages:\r\n",
        "            show_patches(input, prediction, target)\r\n",
        "        \r\n",
        "        # BATCH EVALUATION\r\n",
        "        metrics_batch = calc_metrics(target, prediction, verbose=True)\r\n",
        "        # BAND-WISE EVALUATION\r\n",
        "        metrics['PCC'] += metrics_batch['PCC']\r\n",
        "        metrics['RMSE'] += metrics_batch['RMSE']\r\n",
        "        metrics['PSNR'] += metrics_batch['PSNR']\r\n",
        "        metrics['SSIM'] += metrics_batch['SSIM']\r\n",
        "        # PIXEL-WISE EVALUATION\r\n",
        "        metrics['SAM'] += metrics_batch['SAM']\r\n",
        "        metrics['SID'] += metrics_batch['SID']\r\n",
        "\r\n",
        "\r\n",
        "        '''\r\n",
        "        if saveMetrics != None:\r\n",
        "            metrics = {k: np.mean(m) for k,m in metrics.items()}\r\n",
        "            df = pd.DataFrame({key: pd.Series(value) for key, value in metrics.items()})\r\n",
        "            df.to_csv(os.getcwd() + f\"/drive/My Drive/TFG/Metrics/P2P_metrics/{saveMetrics}.csv\", encoding='utf-8', index=False)\r\n",
        "            break\r\n",
        "        '''\r\n",
        "        \r\n",
        "        # CROP CLASSIFICATION\r\n",
        "        if svc != None:\r\n",
        "            crop = np.array(crop)\r\n",
        "            crop_class, pred_class = svc.test(crop, predictions)\r\n",
        "            print('Accuracy:', sklearn.metrics.accuracy_score(crop_class, pred_class))\r\n",
        "\r\n",
        "    # DATASET EVALUATION\r\n",
        "    metrics = {k: m/5 for k,m in metrics.items()}\r\n",
        "    show_metrics(metrics)\r\n",
        "\r\n",
        "    if saveMetrics != None:\r\n",
        "        df = pd.DataFrame({key: pd.Series(value) for key, value in metrics.items()})\r\n",
        "        df.to_csv(os.getcwd() + f\"/drive/My Drive/TFG/Metrics/P2P_metrics/{saveMetrics}.csv\", encoding='utf-8', index=False)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}