Siamese_Concatenation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "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.7.6"
    },
    "colab": {
      "name": "Siamese Concatenation.ipynb",
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/purbayankar/FeatureFusionUNet/blob/main/Siamese_Concatenation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BxcLflB9uMT8",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "5dc0df7f-ce64-4285-d0aa-1a9bec34782c"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EBK-VDWitCCD",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "2cbd6428-b20f-4f50-ad6e-1848f2caca6b"
      },
      "source": [
        "# Imports\n",
        "\n",
        "# PyTorch\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.autograd import Variable\n",
        "import torchvision.transforms as tr\n",
        "\n",
        "# Models\n",
        "# from unet import Unet\n",
        "# from siamunet_conc import SiamUnet_conc\n",
        "# from siamunet_diff import SiamUnet_diff\n",
        "# from fresunet import FresUNet\n",
        "\n",
        "# Other\n",
        "import os\n",
        "import numpy as np\n",
        "import random\n",
        "from skimage import io\n",
        "from scipy.ndimage import zoom\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "from tqdm import tqdm as tqdm\n",
        "from pandas import read_csv\n",
        "from math import floor, ceil, sqrt, exp\n",
        "from IPython import display\n",
        "import time\n",
        "from itertools import chain\n",
        "import time\n",
        "import warnings\n",
        "from pprint import pprint\n",
        "\n",
        "\n",
        "\n",
        "print('IMPORTS OK')\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "IMPORTS OK\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1RPcV8BktCCW",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "54a0ec4c-0be3-4555-f729-4d1da3991908"
      },
      "source": [
        "# Global Variables' Definitions\n",
        "\n",
        "PATH_TO_DATASET = '/content/drive/My Drive/OneraDataset_Images/'\n",
        "IS_PROTOTYPE = False\n",
        "\n",
        "FP_MODIFIER = 10 # Tuning parameter, use 1 if unsure\n",
        "\n",
        "BATCH_SIZE = 32\n",
        "PATCH_SIDE = 96\n",
        "N_EPOCHS = 50\n",
        "\n",
        "NORMALISE_IMGS = True\n",
        "\n",
        "TRAIN_STRIDE = int(PATCH_SIDE/2) - 1\n",
        "\n",
        "TYPE = 3 # 0-RGB | 1-RGBIr | 2-All bands s.t. resulution <= 20m | 3-All bands\n",
        "\n",
        "LOAD_TRAINED = False\n",
        "\n",
        "DATA_AUG = True\n",
        "\n",
        "\n",
        "print('DEFINITIONS OK')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "DEFINITIONS OK\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hgI-4LV3tCCj",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "9479ef93-b71b-48cd-a8e8-5ef6be976fc6"
      },
      "source": [
        "# Functions\n",
        "\n",
        "def adjust_shape(I, s):\n",
        "    \"\"\"Adjust shape of grayscale image I to s.\"\"\"\n",
        "    \n",
        "    # crop if necesary\n",
        "    I = I[:s[0],:s[1]]\n",
        "    si = I.shape\n",
        "    \n",
        "    # pad if necessary \n",
        "    p0 = max(0,s[0] - si[0])\n",
        "    p1 = max(0,s[1] - si[1])\n",
        "    \n",
        "    return np.pad(I,((0,p0),(0,p1)),'edge')\n",
        "    \n",
        "\n",
        "def read_sentinel_img(path):\n",
        "    \"\"\"Read cropped Sentinel-2 image: RGB bands.\"\"\"\n",
        "    im_name = os.listdir(path)[0][:-7]\n",
        "    r = io.imread(path + im_name + \"B04.tif\")\n",
        "    g = io.imread(path + im_name + \"B03.tif\")\n",
        "    b = io.imread(path + im_name + \"B02.tif\")\n",
        "    \n",
        "    I = np.stack((r,g,b),axis=2).astype('float')\n",
        "    \n",
        "    if NORMALISE_IMGS:\n",
        "        I = (I - I.mean()) / I.std()\n",
        "\n",
        "    return I\n",
        "\n",
        "def read_sentinel_img_4(path):\n",
        "    \"\"\"Read cropped Sentinel-2 image: RGB and NIR bands.\"\"\"\n",
        "    im_name = os.listdir(path)[0][:-7]\n",
        "    r = io.imread(path + im_name + \"B04.tif\")\n",
        "    g = io.imread(path + im_name + \"B03.tif\")\n",
        "    b = io.imread(path + im_name + \"B02.tif\")\n",
        "    nir = io.imread(path + im_name + \"B08.tif\")\n",
        "    \n",
        "    I = np.stack((r,g,b,nir),axis=2).astype('float')\n",
        "    \n",
        "    if NORMALISE_IMGS:\n",
        "        I = (I - I.mean()) / I.std()\n",
        "\n",
        "    return I\n",
        "\n",
        "def read_sentinel_img_leq20(path):\n",
        "    \"\"\"Read cropped Sentinel-2 image: bands with resolution less than or equals to 20m.\"\"\"\n",
        "    im_name = os.listdir(path)[0][:-7]\n",
        "    \n",
        "    r = io.imread(path + im_name + \"B04.tif\")\n",
        "    s = r.shape\n",
        "    g = io.imread(path + im_name + \"B03.tif\")\n",
        "    b = io.imread(path + im_name + \"B02.tif\")\n",
        "    nir = io.imread(path + im_name + \"B08.tif\")\n",
        "    \n",
        "    ir1 = adjust_shape(zoom(io.imread(path + im_name + \"B05.tif\"),2),s)\n",
        "    ir2 = adjust_shape(zoom(io.imread(path + im_name + \"B06.tif\"),2),s)\n",
        "    ir3 = adjust_shape(zoom(io.imread(path + im_name + \"B07.tif\"),2),s)\n",
        "    nir2 = adjust_shape(zoom(io.imread(path + im_name + \"B8A.tif\"),2),s)\n",
        "    swir2 = adjust_shape(zoom(io.imread(path + im_name + \"B11.tif\"),2),s)\n",
        "    swir3 = adjust_shape(zoom(io.imread(path + im_name + \"B12.tif\"),2),s)\n",
        "    \n",
        "    I = np.stack((r,g,b,nir,ir1,ir2,ir3,nir2,swir2,swir3),axis=2).astype('float')\n",
        "    \n",
        "    if NORMALISE_IMGS:\n",
        "        I = (I - I.mean()) / I.std()\n",
        "\n",
        "    return I\n",
        "\n",
        "def read_sentinel_img_leq60(path):\n",
        "    \"\"\"Read cropped Sentinel-2 image: all bands.\"\"\"\n",
        "    im_name = os.listdir(path)[0][:-7]\n",
        "    \n",
        "    r = io.imread(path + im_name + \"B04.tif\")\n",
        "    s = r.shape\n",
        "    g = io.imread(path + im_name + \"B03.tif\")\n",
        "    b = io.imread(path + im_name + \"B02.tif\")\n",
        "    nir = io.imread(path + im_name + \"B08.tif\")\n",
        "    \n",
        "    ir1 = adjust_shape(zoom(io.imread(path + im_name + \"B05.tif\"),2),s)\n",
        "    ir2 = adjust_shape(zoom(io.imread(path + im_name + \"B06.tif\"),2),s)\n",
        "    ir3 = adjust_shape(zoom(io.imread(path + im_name + \"B07.tif\"),2),s)\n",
        "    nir2 = adjust_shape(zoom(io.imread(path + im_name + \"B8A.tif\"),2),s)\n",
        "    swir2 = adjust_shape(zoom(io.imread(path + im_name + \"B11.tif\"),2),s)\n",
        "    swir3 = adjust_shape(zoom(io.imread(path + im_name + \"B12.tif\"),2),s)\n",
        "    \n",
        "    uv = adjust_shape(zoom(io.imread(path + im_name + \"B01.tif\"),6),s)\n",
        "    wv = adjust_shape(zoom(io.imread(path + im_name + \"B09.tif\"),6),s)\n",
        "    swirc = adjust_shape(zoom(io.imread(path + im_name + \"B10.tif\"),6),s)\n",
        "    \n",
        "    I = np.stack((r,g,b,nir,ir1,ir2,ir3,nir2,swir2,swir3,uv,wv,swirc),axis=2).astype('float')\n",
        "    \n",
        "    if NORMALISE_IMGS:\n",
        "        I = (I - I.mean()) / I.std()\n",
        "\n",
        "    return I\n",
        "\n",
        "def read_sentinel_img_trio(path):\n",
        "    \"\"\"Read cropped Sentinel-2 image pair and change map.\"\"\"\n",
        "#     read images\n",
        "    if TYPE == 0:\n",
        "        I1 = read_sentinel_img(path + '/imgs_1/')\n",
        "        I2 = read_sentinel_img(path + '/imgs_2/')\n",
        "    elif TYPE == 1:\n",
        "        I1 = read_sentinel_img_4(path + '/imgs_1/')\n",
        "        I2 = read_sentinel_img_4(path + '/imgs_2/')\n",
        "    elif TYPE == 2:\n",
        "        I1 = read_sentinel_img_leq20(path + '/imgs_1/')\n",
        "        I2 = read_sentinel_img_leq20(path + '/imgs_2/')\n",
        "    elif TYPE == 3:\n",
        "        I1 = read_sentinel_img_leq60(path + '/imgs_1/')\n",
        "        I2 = read_sentinel_img_leq60(path + '/imgs_2/')\n",
        "        \n",
        "    cm = io.imread(path + '/cm/cm.png', as_gray=True) != 0\n",
        "    \n",
        "    # crop if necessary\n",
        "    s1 = I1.shape\n",
        "    s2 = I2.shape\n",
        "    I2 = np.pad(I2,((0, s1[0] - s2[0]), (0, s1[1] - s2[1]), (0,0)),'edge')\n",
        "    \n",
        "    \n",
        "    return I1, I2, cm\n",
        "\n",
        "\n",
        "\n",
        "def reshape_for_torch(I):\n",
        "    \"\"\"Transpose image for PyTorch coordinates.\"\"\"\n",
        "#     out = np.swapaxes(I,1,2)\n",
        "#     out = np.swapaxes(out,0,1)\n",
        "#     out = out[np.newaxis,:]\n",
        "    out = I.transpose((2, 0, 1))\n",
        "    return torch.from_numpy(out)\n",
        "\n",
        "\n",
        "\n",
        "class ChangeDetectionDataset(Dataset):\n",
        "    \"\"\"Change Detection dataset class, used for both training and test data.\"\"\"\n",
        "\n",
        "    def __init__(self, path, train = True, patch_side = 96, stride = None, use_all_bands = False, transform=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            csv_file (string): Path to the csv file with annotations.\n",
        "            root_dir (string): Directory with all the images.\n",
        "            transform (callable, optional): Optional transform to be applied\n",
        "                on a sample.\n",
        "        \"\"\"\n",
        "        \n",
        "        # basics\n",
        "        self.transform = transform\n",
        "        self.path = path\n",
        "        self.patch_side = patch_side\n",
        "        if not stride:\n",
        "            self.stride = 1\n",
        "        else:\n",
        "            self.stride = stride\n",
        "        \n",
        "        if train:\n",
        "            fname = 'train.txt'\n",
        "        else:\n",
        "            fname = 'test.txt'\n",
        "        \n",
        "#         print(path + fname)\n",
        "        self.names = read_csv(path + fname).columns\n",
        "        self.n_imgs = self.names.shape[0]\n",
        "        \n",
        "        n_pix = 0\n",
        "        true_pix = 0\n",
        "        \n",
        "        \n",
        "        # load images\n",
        "        self.imgs_1 = {}\n",
        "        self.imgs_2 = {}\n",
        "        self.change_maps = {}\n",
        "        self.n_patches_per_image = {}\n",
        "        self.n_patches = 0\n",
        "        self.patch_coords = []\n",
        "        for im_name in tqdm(self.names):\n",
        "            # load and store each image\n",
        "            I1, I2, cm = read_sentinel_img_trio(self.path + im_name)\n",
        "            self.imgs_1[im_name] = reshape_for_torch(I1)\n",
        "            self.imgs_2[im_name] = reshape_for_torch(I2)\n",
        "            self.change_maps[im_name] = cm\n",
        "            \n",
        "            s = cm.shape\n",
        "            n_pix += np.prod(s)\n",
        "            true_pix += cm.sum()\n",
        "            \n",
        "            # calculate the number of patches\n",
        "            s = self.imgs_1[im_name].shape\n",
        "            n1 = ceil((s[1] - self.patch_side + 1) / self.stride)\n",
        "            n2 = ceil((s[2] - self.patch_side + 1) / self.stride)\n",
        "            n_patches_i = n1 * n2\n",
        "            self.n_patches_per_image[im_name] = n_patches_i\n",
        "            self.n_patches += n_patches_i\n",
        "            \n",
        "            # generate path coordinates\n",
        "            for i in range(n1):\n",
        "                for j in range(n2):\n",
        "                    # coordinates in (x1, x2, y1, y2)\n",
        "                    current_patch_coords = (im_name, \n",
        "                                    [self.stride*i, self.stride*i + self.patch_side, self.stride*j, self.stride*j + self.patch_side],\n",
        "                                    [self.stride*(i + 1), self.stride*(j + 1)])\n",
        "                    self.patch_coords.append(current_patch_coords)\n",
        "                    \n",
        "        self.weights = [ FP_MODIFIER * 2 * true_pix / n_pix, 2 * (n_pix - true_pix) / n_pix]\n",
        "        \n",
        "        \n",
        "\n",
        "    def get_img(self, im_name):\n",
        "        return self.imgs_1[im_name], self.imgs_2[im_name], self.change_maps[im_name]\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.n_patches\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        current_patch_coords = self.patch_coords[idx]\n",
        "        im_name = current_patch_coords[0]\n",
        "        limits = current_patch_coords[1]\n",
        "        centre = current_patch_coords[2]\n",
        "        \n",
        "        I1 = self.imgs_1[im_name][:, limits[0]:limits[1], limits[2]:limits[3]]\n",
        "        I2 = self.imgs_2[im_name][:, limits[0]:limits[1], limits[2]:limits[3]]\n",
        "        \n",
        "        label = self.change_maps[im_name][limits[0]:limits[1], limits[2]:limits[3]]\n",
        "        label = torch.from_numpy(1*np.array(label)).float()\n",
        "        \n",
        "        sample = {'I1': I1, 'I2': I2, 'label': label}\n",
        "        \n",
        "        if self.transform:\n",
        "            sample = self.transform(sample)\n",
        "\n",
        "        return sample\n",
        "\n",
        "\n",
        "\n",
        "class RandomFlip(object):\n",
        "    \"\"\"Flip randomly the images in a sample.\"\"\"\n",
        "\n",
        "#     def __init__(self):\n",
        "#         return\n",
        "\n",
        "    def __call__(self, sample):\n",
        "        I1, I2, label = sample['I1'], sample['I2'], sample['label']\n",
        "        \n",
        "        if random.random() > 0.5:\n",
        "            I1 =  I1.numpy()[:,:,::-1].copy()\n",
        "            I1 = torch.from_numpy(I1)\n",
        "            I2 =  I2.numpy()[:,:,::-1].copy()\n",
        "            I2 = torch.from_numpy(I2)\n",
        "            label =  label.numpy()[:,::-1].copy()\n",
        "            label = torch.from_numpy(label)\n",
        "\n",
        "        return {'I1': I1, 'I2': I2, 'label': label}\n",
        "\n",
        "\n",
        "\n",
        "class RandomRot(object):\n",
        "    \"\"\"Rotate randomly the images in a sample.\"\"\"\n",
        "\n",
        "#     def __init__(self):\n",
        "#         return\n",
        "\n",
        "    def __call__(self, sample):\n",
        "        I1, I2, label = sample['I1'], sample['I2'], sample['label']\n",
        "        \n",
        "        n = random.randint(0, 3)\n",
        "        if n:\n",
        "            I1 =  sample['I1'].numpy()\n",
        "            I1 = np.rot90(I1, n, axes=(1, 2)).copy()\n",
        "            I1 = torch.from_numpy(I1)\n",
        "            I2 =  sample['I2'].numpy()\n",
        "            I2 = np.rot90(I2, n, axes=(1, 2)).copy()\n",
        "            I2 = torch.from_numpy(I2)\n",
        "            label =  sample['label'].numpy()\n",
        "            label = np.rot90(label, n, axes=(0, 1)).copy()\n",
        "            label = torch.from_numpy(label)\n",
        "\n",
        "        return {'I1': I1, 'I2': I2, 'label': label}\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "print('UTILS OK')\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "UTILS OK\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V5l3hr1AtCCt"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_rDHwU5gtCC1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 87
        },
        "outputId": "71d116e7-87f0-40cb-d176-10b0d6e29bd3"
      },
      "source": [
        "# Dataset\n",
        "\n",
        "\n",
        "if DATA_AUG:\n",
        "    data_transform = tr.Compose([RandomFlip(), RandomRot()])\n",
        "else:\n",
        "    data_transform = None\n",
        "\n",
        "\n",
        "        \n",
        "\n",
        "train_dataset = ChangeDetectionDataset(PATH_TO_DATASET, train = True, stride = TRAIN_STRIDE, transform=data_transform)\n",
        "weights = torch.FloatTensor(train_dataset.weights).cuda()\n",
        "print(weights)\n",
        "train_loader = DataLoader(train_dataset, batch_size = BATCH_SIZE, shuffle = True, num_workers = 4)\n",
        "\n",
        "test_dataset = ChangeDetectionDataset(PATH_TO_DATASET, train = False, stride = TRAIN_STRIDE)\n",
        "test_loader = DataLoader(test_dataset, batch_size = BATCH_SIZE, shuffle = True, num_workers = 4)\n",
        "\n",
        "\n",
        "print('DATASETS OK')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 14/14 [02:27<00:00, 10.53s/it]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "tensor([0.4595, 1.9540], device='cuda:0')\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 10/10 [01:37<00:00,  9.79s/it]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "DATASETS OK\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7s94wymZtCDD"
      },
      "source": [
        "# print(weights)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HBSQAxAutlpM"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.nn.modules.padding import ReplicationPad2d\n",
        "\n",
        "class SiamUnet_conc(nn.Module):\n",
        "    \"\"\"SiamUnet_conc segmentation network.\"\"\"\n",
        "\n",
        "    def __init__(self, input_nbr, label_nbr):\n",
        "        super(SiamUnet_conc, self).__init__()\n",
        "\n",
        "        self.input_nbr = input_nbr\n",
        "\n",
        "        self.conv11 = nn.Conv2d(input_nbr, 16, kernel_size=3, padding=1)\n",
        "        self.bn11 = nn.BatchNorm2d(16)\n",
        "        self.do11 = nn.Dropout2d(p=0.2)\n",
        "        self.conv12 = nn.Conv2d(16, 16, kernel_size=3, padding=1)\n",
        "        self.bn12 = nn.BatchNorm2d(16)\n",
        "        self.do12 = nn.Dropout2d(p=0.2)\n",
        "\n",
        "        self.conv21 = nn.Conv2d(16, 32, kernel_size=3, padding=1)\n",
        "        self.bn21 = nn.BatchNorm2d(32)\n",
        "        self.do21 = nn.Dropout2d(p=0.2)\n",
        "        self.conv22 = nn.Conv2d(32, 32, kernel_size=3, padding=1)\n",
        "        self.bn22 = nn.BatchNorm2d(32)\n",
        "        self.do22 = nn.Dropout2d(p=0.2)\n",
        "\n",
        "        self.conv31 = nn.Conv2d(32, 64, kernel_size=3, padding=1)\n",
        "        self.bn31 = nn.BatchNorm2d(64)\n",
        "        self.do31 = nn.Dropout2d(p=0.2)\n",
        "        self.conv32 = nn.Conv2d(64, 64, kernel_size=3, padding=1)\n",
        "        self.bn32 = nn.BatchNorm2d(64)\n",
        "        self.do32 = nn.Dropout2d(p=0.2)\n",
        "        self.conv33 = nn.Conv2d(64, 64, kernel_size=3, padding=1)\n",
        "        self.bn33 = nn.BatchNorm2d(64)\n",
        "        self.do33 = nn.Dropout2d(p=0.2)\n",
        "\n",
        "        self.conv41 = nn.Conv2d(64, 128, kernel_size=3, padding=1)\n",
        "        self.bn41 = nn.BatchNorm2d(128)\n",
        "        self.do41 = nn.Dropout2d(p=0.2)\n",
        "        self.conv42 = nn.Conv2d(128, 128, kernel_size=3, padding=1)\n",
        "        self.bn42 = nn.BatchNorm2d(128)\n",
        "        self.do42 = nn.Dropout2d(p=0.2)\n",
        "        self.conv43 = nn.Conv2d(128, 128, kernel_size=3, padding=1)\n",
        "        self.bn43 = nn.BatchNorm2d(128)\n",
        "        self.do43 = nn.Dropout2d(p=0.2)\n",
        "\n",
        "        self.upconv4 = nn.ConvTranspose2d(128, 128, kernel_size=3, padding=1, stride=2, output_padding=1)\n",
        "\n",
        "        self.conv43d = nn.ConvTranspose2d(384, 128, kernel_size=3, padding=1)\n",
        "        self.bn43d = nn.BatchNorm2d(128)\n",
        "        self.do43d = nn.Dropout2d(p=0.2)\n",
        "        self.conv42d = nn.ConvTranspose2d(128, 128, kernel_size=3, padding=1)\n",
        "        self.bn42d = nn.BatchNorm2d(128)\n",
        "        self.do42d = nn.Dropout2d(p=0.2)\n",
        "        self.conv41d = nn.ConvTranspose2d(128, 64, kernel_size=3, padding=1)\n",
        "        self.bn41d = nn.BatchNorm2d(64)\n",
        "        self.do41d = nn.Dropout2d(p=0.2)\n",
        "\n",
        "        self.upconv3 = nn.ConvTranspose2d(64, 64, kernel_size=3, padding=1, stride=2, output_padding=1)\n",
        "\n",
        "        self.conv33d = nn.ConvTranspose2d(192, 64, kernel_size=3, padding=1)\n",
        "        self.bn33d = nn.BatchNorm2d(64)\n",
        "        self.do33d = nn.Dropout2d(p=0.2)\n",
        "        self.conv32d = nn.ConvTranspose2d(64, 64, kernel_size=3, padding=1)\n",
        "        self.bn32d = nn.BatchNorm2d(64)\n",
        "        self.do32d = nn.Dropout2d(p=0.2)\n",
        "        self.conv31d = nn.ConvTranspose2d(64, 32, kernel_size=3, padding=1)\n",
        "        self.bn31d = nn.BatchNorm2d(32)\n",
        "        self.do31d = nn.Dropout2d(p=0.2)\n",
        "\n",
        "        self.upconv2 = nn.ConvTranspose2d(32, 32, kernel_size=3, padding=1, stride=2, output_padding=1)\n",
        "\n",
        "        self.conv22d = nn.ConvTranspose2d(96, 32, kernel_size=3, padding=1)\n",
        "        self.bn22d = nn.BatchNorm2d(32)\n",
        "        self.do22d = nn.Dropout2d(p=0.2)\n",
        "        self.conv21d = nn.ConvTranspose2d(32, 16, kernel_size=3, padding=1)\n",
        "        self.bn21d = nn.BatchNorm2d(16)\n",
        "        self.do21d = nn.Dropout2d(p=0.2)\n",
        "\n",
        "        self.upconv1 = nn.ConvTranspose2d(16, 16, kernel_size=3, padding=1, stride=2, output_padding=1)\n",
        "\n",
        "        self.conv12d = nn.ConvTranspose2d(48, 16, kernel_size=3, padding=1)\n",
        "        self.bn12d = nn.BatchNorm2d(16)\n",
        "        self.do12d = nn.Dropout2d(p=0.2)\n",
        "        self.conv11d = nn.ConvTranspose2d(16, label_nbr, kernel_size=3, padding=1)\n",
        "\n",
        "        self.sm = nn.LogSoftmax(dim=1)\n",
        "\n",
        "    def forward(self, x1, x2):\n",
        "\n",
        "        \"\"\"Forward method.\"\"\"\n",
        "        # Stage 1\n",
        "        x11 = self.do11(F.relu(self.bn11(self.conv11(x1))))\n",
        "        x12_1 = self.do12(F.relu(self.bn12(self.conv12(x11))))\n",
        "        x1p = F.max_pool2d(x12_1, kernel_size=2, stride=2)\n",
        "\n",
        "\n",
        "        # Stage 2\n",
        "        x21 = self.do21(F.relu(self.bn21(self.conv21(x1p))))\n",
        "        x22_1 = self.do22(F.relu(self.bn22(self.conv22(x21))))\n",
        "        x2p = F.max_pool2d(x22_1, kernel_size=2, stride=2)\n",
        "\n",
        "        # Stage 3\n",
        "        x31 = self.do31(F.relu(self.bn31(self.conv31(x2p))))\n",
        "        x32 = self.do32(F.relu(self.bn32(self.conv32(x31))))\n",
        "        x33_1 = self.do33(F.relu(self.bn33(self.conv33(x32))))\n",
        "        x3p = F.max_pool2d(x33_1, kernel_size=2, stride=2)\n",
        "\n",
        "        # Stage 4\n",
        "        x41 = self.do41(F.relu(self.bn41(self.conv41(x3p))))\n",
        "        x42 = self.do42(F.relu(self.bn42(self.conv42(x41))))\n",
        "        x43_1 = self.do43(F.relu(self.bn43(self.conv43(x42))))\n",
        "        x4p = F.max_pool2d(x43_1, kernel_size=2, stride=2)\n",
        "\n",
        "\n",
        "        ####################################################\n",
        "        # Stage 1\n",
        "        x11 = self.do11(F.relu(self.bn11(self.conv11(x2))))\n",
        "        x12_2 = self.do12(F.relu(self.bn12(self.conv12(x11))))\n",
        "        x1p = F.max_pool2d(x12_2, kernel_size=2, stride=2)\n",
        "\n",
        "        # Stage 2\n",
        "        x21 = self.do21(F.relu(self.bn21(self.conv21(x1p))))\n",
        "        x22_2 = self.do22(F.relu(self.bn22(self.conv22(x21))))\n",
        "        x2p = F.max_pool2d(x22_2, kernel_size=2, stride=2)\n",
        "\n",
        "        # Stage 3\n",
        "        x31 = self.do31(F.relu(self.bn31(self.conv31(x2p))))\n",
        "        x32 = self.do32(F.relu(self.bn32(self.conv32(x31))))\n",
        "        x33_2 = self.do33(F.relu(self.bn33(self.conv33(x32))))\n",
        "        x3p = F.max_pool2d(x33_2, kernel_size=2, stride=2)\n",
        "\n",
        "        # Stage 4\n",
        "        x41 = self.do41(F.relu(self.bn41(self.conv41(x3p))))\n",
        "        x42 = self.do42(F.relu(self.bn42(self.conv42(x41))))\n",
        "        x43_2 = self.do43(F.relu(self.bn43(self.conv43(x42))))\n",
        "        x4p = F.max_pool2d(x43_2, kernel_size=2, stride=2)\n",
        "\n",
        "\n",
        "        ####################################################\n",
        "        # Stage 4d\n",
        "        x4d = self.upconv4(x4p)\n",
        "        pad4 = ReplicationPad2d((0, x43_1.size(3) - x4d.size(3), 0, x43_1.size(2) - x4d.size(2)))\n",
        "        x4d = torch.cat((pad4(x4d), x43_1, x43_2), 1)\n",
        "        x43d = self.do43d(F.relu(self.bn43d(self.conv43d(x4d))))\n",
        "        x42d = self.do42d(F.relu(self.bn42d(self.conv42d(x43d))))\n",
        "        x41d = self.do41d(F.relu(self.bn41d(self.conv41d(x42d))))\n",
        "\n",
        "        # Stage 3d\n",
        "        x3d = self.upconv3(x41d)\n",
        "        pad3 = ReplicationPad2d((0, x33_1.size(3) - x3d.size(3), 0, x33_1.size(2) - x3d.size(2)))\n",
        "        x3d = torch.cat((pad3(x3d), x33_1, x33_2), 1)\n",
        "        x33d = self.do33d(F.relu(self.bn33d(self.conv33d(x3d))))\n",
        "        x32d = self.do32d(F.relu(self.bn32d(self.conv32d(x33d))))\n",
        "        x31d = self.do31d(F.relu(self.bn31d(self.conv31d(x32d))))\n",
        "\n",
        "        # Stage 2d\n",
        "        x2d = self.upconv2(x31d)\n",
        "        pad2 = ReplicationPad2d((0, x22_1.size(3) - x2d.size(3), 0, x22_1.size(2) - x2d.size(2)))\n",
        "        x2d = torch.cat((pad2(x2d), x22_1, x22_2), 1)\n",
        "        x22d = self.do22d(F.relu(self.bn22d(self.conv22d(x2d))))\n",
        "        x21d = self.do21d(F.relu(self.bn21d(self.conv21d(x22d))))\n",
        "\n",
        "        # Stage 1d\n",
        "        x1d = self.upconv1(x21d)\n",
        "        pad1 = ReplicationPad2d((0, x12_1.size(3) - x1d.size(3), 0, x12_1.size(2) - x1d.size(2)))\n",
        "        x1d = torch.cat((pad1(x1d), x12_1, x12_2), 1)\n",
        "        x12d = self.do12d(F.relu(self.bn12d(self.conv12d(x1d))))\n",
        "        x11d = self.conv11d(x12d)\n",
        "\n",
        "        return self.sm(x11d)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-_-6nJdHtCDL",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "83431411-36a2-4adc-e492-5f464d1b7d05"
      },
      "source": [
        "# 0-RGB | 1-RGBIr | 2-All bands s.t. resulution <= 20m | 3-All bands\n",
        "\n",
        "if TYPE == 0:\n",
        "    net, net_name = Unet(2*3, 2), 'FC-EF'\n",
        "#     net, net_name = SiamUnet_conc(3, 2), 'FC-Siam-conc'\n",
        "#     net, net_name = SiamUnet_diff(3, 2), 'FC-Siam-diff'\n",
        "    # net, net_name = FresUNet(2*3, 2), 'FresUNet'\n",
        "elif TYPE == 1:\n",
        "    net, net_name = Unet(2*4, 2), 'FC-EF'\n",
        "#     net, net_name = SiamUnet_conc(4, 2), 'FC-Siam-conc'\n",
        "#     net, net_name = SiamUnet_diff(4, 2), 'FC-Siam-diff'\n",
        "    # net, net_name = FresUNet(2*4, 2), 'FresUNet'\n",
        "elif TYPE == 2:\n",
        "    net, net_name = Unet(2*10, 2), 'FC-EF'\n",
        "#     net, net_name = SiamUnet_conc(10, 2), 'FC-Siam-conc'\n",
        "#     net, net_name = SiamUnet_diff(10, 2), 'FC-Siam-diff'\n",
        "#     net, net_name = FresUNet(2*10, 2), 'FresUNet'\n",
        "elif TYPE == 3:\n",
        "#     net, net_name = Unet(2*13, 2), 'FC-EF'\n",
        "    net, net_name = SiamUnet_conc(13, 2), 'FC-Siam-conc'\n",
        "#     net, net_name = SiamUnet_diff(13, 2), 'FC-Siam-diff'\n",
        "#     net, net_name = FresUNet(2*13, 2), 'FresUNet'\n",
        "\n",
        "\n",
        "net.cuda()\n",
        "\n",
        "criterion = nn.NLLLoss(weight=weights) # to be used with logsoftmax output\n",
        "from torchsummary import summary\n",
        "\n",
        "summary(net,input_size=[(13,256,256),(13,256,256)])\n",
        "\n",
        "print('NETWORK OK')\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "----------------------------------------------------------------\n",
            "        Layer (type)               Output Shape         Param #\n",
            "================================================================\n",
            "            Conv2d-1         [-1, 32, 256, 256]           7,520\n",
            "       BatchNorm2d-2         [-1, 32, 256, 256]              64\n",
            "              ReLU-3         [-1, 32, 256, 256]               0\n",
            "            Conv2d-4         [-1, 32, 256, 256]           9,248\n",
            "       BatchNorm2d-5         [-1, 32, 256, 256]              64\n",
            "              ReLU-6         [-1, 32, 256, 256]               0\n",
            " conv_block_nested-7         [-1, 32, 256, 256]               0\n",
            "         MaxPool2d-8         [-1, 32, 128, 128]               0\n",
            "            Conv2d-9         [-1, 64, 128, 128]          18,496\n",
            "      BatchNorm2d-10         [-1, 64, 128, 128]             128\n",
            "             ReLU-11         [-1, 64, 128, 128]               0\n",
            "           Conv2d-12         [-1, 64, 128, 128]          36,928\n",
            "      BatchNorm2d-13         [-1, 64, 128, 128]             128\n",
            "             ReLU-14         [-1, 64, 128, 128]               0\n",
            "conv_block_nested-15         [-1, 64, 128, 128]               0\n",
            "        MaxPool2d-16           [-1, 64, 64, 64]               0\n",
            "           Conv2d-17          [-1, 128, 64, 64]          73,856\n",
            "      BatchNorm2d-18          [-1, 128, 64, 64]             256\n",
            "             ReLU-19          [-1, 128, 64, 64]               0\n",
            "           Conv2d-20          [-1, 128, 64, 64]         147,584\n",
            "      BatchNorm2d-21          [-1, 128, 64, 64]             256\n",
            "             ReLU-22          [-1, 128, 64, 64]               0\n",
            "conv_block_nested-23          [-1, 128, 64, 64]               0\n",
            "        MaxPool2d-24          [-1, 128, 32, 32]               0\n",
            "           Conv2d-25          [-1, 256, 32, 32]         295,168\n",
            "      BatchNorm2d-26          [-1, 256, 32, 32]             512\n",
            "             ReLU-27          [-1, 256, 32, 32]               0\n",
            "           Conv2d-28          [-1, 256, 32, 32]         590,080\n",
            "      BatchNorm2d-29          [-1, 256, 32, 32]             512\n",
            "             ReLU-30          [-1, 256, 32, 32]               0\n",
            "conv_block_nested-31          [-1, 256, 32, 32]               0\n",
            "        MaxPool2d-32          [-1, 256, 16, 16]               0\n",
            "           Conv2d-33          [-1, 512, 16, 16]       1,180,160\n",
            "      BatchNorm2d-34          [-1, 512, 16, 16]           1,024\n",
            "             ReLU-35          [-1, 512, 16, 16]               0\n",
            "           Conv2d-36          [-1, 512, 16, 16]       2,359,808\n",
            "      BatchNorm2d-37          [-1, 512, 16, 16]           1,024\n",
            "             ReLU-38          [-1, 512, 16, 16]               0\n",
            "conv_block_nested-39          [-1, 512, 16, 16]               0\n",
            "         Upsample-40          [-1, 512, 32, 32]               0\n",
            "           Conv2d-41          [-1, 256, 32, 32]       1,769,728\n",
            "      BatchNorm2d-42          [-1, 256, 32, 32]             512\n",
            "             ReLU-43          [-1, 256, 32, 32]               0\n",
            "           Conv2d-44          [-1, 256, 32, 32]         590,080\n",
            "      BatchNorm2d-45          [-1, 256, 32, 32]             512\n",
            "             ReLU-46          [-1, 256, 32, 32]               0\n",
            "conv_block_nested-47          [-1, 256, 32, 32]               0\n",
            "         Upsample-48          [-1, 256, 64, 64]               0\n",
            "           Conv2d-49          [-1, 128, 64, 64]         442,496\n",
            "      BatchNorm2d-50          [-1, 128, 64, 64]             256\n",
            "             ReLU-51          [-1, 128, 64, 64]               0\n",
            "           Conv2d-52          [-1, 128, 64, 64]         147,584\n",
            "      BatchNorm2d-53          [-1, 128, 64, 64]             256\n",
            "             ReLU-54          [-1, 128, 64, 64]               0\n",
            "conv_block_nested-55          [-1, 128, 64, 64]               0\n",
            "         Upsample-56        [-1, 128, 128, 128]               0\n",
            "           Conv2d-57         [-1, 64, 128, 128]         110,656\n",
            "      BatchNorm2d-58         [-1, 64, 128, 128]             128\n",
            "             ReLU-59         [-1, 64, 128, 128]               0\n",
            "           Conv2d-60         [-1, 64, 128, 128]          36,928\n",
            "      BatchNorm2d-61         [-1, 64, 128, 128]             128\n",
            "             ReLU-62         [-1, 64, 128, 128]               0\n",
            "conv_block_nested-63         [-1, 64, 128, 128]               0\n",
            "         Upsample-64         [-1, 64, 256, 256]               0\n",
            "           Conv2d-65         [-1, 32, 256, 256]          27,680\n",
            "      BatchNorm2d-66         [-1, 32, 256, 256]              64\n",
            "             ReLU-67         [-1, 32, 256, 256]               0\n",
            "           Conv2d-68         [-1, 32, 256, 256]           9,248\n",
            "      BatchNorm2d-69         [-1, 32, 256, 256]              64\n",
            "             ReLU-70         [-1, 32, 256, 256]               0\n",
            "conv_block_nested-71         [-1, 32, 256, 256]               0\n",
            "           Conv2d-72          [-1, 2, 256, 256]              66\n",
            "================================================================\n",
            "Total params: 7,859,202\n",
            "Trainable params: 7,859,202\n",
            "Non-trainable params: 0\n",
            "----------------------------------------------------------------\n",
            "Input size (MB): 2768896.00\n",
            "Forward/backward pass size (MB): 495.50\n",
            "Params size (MB): 29.98\n",
            "Estimated Total Size (MB): 2769421.48\n",
            "----------------------------------------------------------------\n",
            "NETWORK OK\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PeLsY037tCDT",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "d368106d-3080-4169-cd12-486505326a8e"
      },
      "source": [
        "def count_parameters(model):\n",
        "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "\n",
        "print('Number of trainable parameters:', count_parameters(net))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Number of trainable parameters: 7859202\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p7dfDWGCtCDb"
      },
      "source": [
        "# net.load_state_dict(torch.load('net-best_epoch-1_fm-0.7394933126157746.pth.tar'))\n",
        "\n",
        "def train(n_epochs = N_EPOCHS, save = True):\n",
        "    t = np.linspace(1, n_epochs, n_epochs)\n",
        "    \n",
        "    epoch_train_loss = 0 * t\n",
        "    epoch_train_accuracy = 0 * t\n",
        "    epoch_train_change_accuracy = 0 * t\n",
        "    epoch_train_nochange_accuracy = 0 * t\n",
        "    epoch_train_precision = 0 * t\n",
        "    epoch_train_recall = 0 * t\n",
        "    epoch_train_Fmeasure = 0 * t\n",
        "    epoch_test_loss = 0 * t\n",
        "    epoch_test_accuracy = 0 * t\n",
        "    epoch_test_change_accuracy = 0 * t\n",
        "    epoch_test_nochange_accuracy = 0 * t\n",
        "    epoch_test_precision = 0 * t\n",
        "    epoch_test_recall = 0 * t\n",
        "    epoch_test_Fmeasure = 0 * t\n",
        "    \n",
        "#     mean_acc = 0\n",
        "#     best_mean_acc = 0\n",
        "    fm = 0\n",
        "    best_fm = 0\n",
        "    \n",
        "    lss = 1000\n",
        "    best_lss = 1000\n",
        "    \n",
        "    plt.figure(num=1)\n",
        "    plt.figure(num=2)\n",
        "    plt.figure(num=3)\n",
        "    \n",
        "    \n",
        "    optimizer = torch.optim.Adam(net.parameters(), weight_decay=1e-4)\n",
        "#     optimizer = torch.optim.Adam(net.parameters(), lr=0.0005)\n",
        "        \n",
        "    \n",
        "    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, 0.95)\n",
        "    \n",
        "    \n",
        "    for epoch_index in tqdm(range(n_epochs)):\n",
        "        net.train()\n",
        "        print('Epoch: ' + str(epoch_index + 1) + ' of ' + str(N_EPOCHS))\n",
        "\n",
        "        tot_count = 0\n",
        "        tot_loss = 0\n",
        "        tot_accurate = 0\n",
        "        class_correct = list(0. for i in range(2))\n",
        "        class_total = list(0. for i in range(2))\n",
        "#         for batch_index, batch in enumerate(tqdm(data_loader)):\n",
        "        for batch in train_loader:\n",
        "            I1 = Variable(batch['I1'].float().cuda())\n",
        "            I2 = Variable(batch['I2'].float().cuda())\n",
        "            label = torch.squeeze(Variable(batch['label'].cuda()))\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            output = net(I1, I2)\n",
        "            loss = criterion(output, label.long())\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            \n",
        "        scheduler.step()\n",
        "\n",
        "\n",
        "        epoch_train_loss[epoch_index], epoch_train_accuracy[epoch_index], cl_acc, pr_rec = test(train_dataset)\n",
        "        epoch_train_nochange_accuracy[epoch_index] = cl_acc[0]\n",
        "        epoch_train_change_accuracy[epoch_index] = cl_acc[1]\n",
        "        epoch_train_precision[epoch_index] = pr_rec[0]\n",
        "        epoch_train_recall[epoch_index] = pr_rec[1]\n",
        "        epoch_train_Fmeasure[epoch_index] = pr_rec[2]\n",
        "        \n",
        "#         epoch_test_loss[epoch_index], epoch_test_accuracy[epoch_index], cl_acc, pr_rec = test(test_dataset)\n",
        "        epoch_test_loss[epoch_index], epoch_test_accuracy[epoch_index], cl_acc, pr_rec = test(test_dataset)\n",
        "        epoch_test_nochange_accuracy[epoch_index] = cl_acc[0]\n",
        "        epoch_test_change_accuracy[epoch_index] = cl_acc[1]\n",
        "        epoch_test_precision[epoch_index] = pr_rec[0]\n",
        "        epoch_test_recall[epoch_index] = pr_rec[1]\n",
        "        epoch_test_Fmeasure[epoch_index] = pr_rec[2]\n",
        "\n",
        "        plt.figure(num=1)\n",
        "        plt.clf()\n",
        "        l1_1, = plt.plot(t[:epoch_index + 1], epoch_train_loss[:epoch_index + 1], label='Train loss')\n",
        "        l1_2, = plt.plot(t[:epoch_index + 1], epoch_test_loss[:epoch_index + 1], label='Test loss')\n",
        "        plt.legend(handles=[l1_1, l1_2])\n",
        "        plt.grid()\n",
        "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
        "        plt.gcf().gca().set_xlim(left = 0)\n",
        "        plt.title('Loss')\n",
        "        display.clear_output(wait=True)\n",
        "        display.display(plt.gcf())\n",
        "\n",
        "        plt.figure(num=2)\n",
        "        plt.clf()\n",
        "        l2_1, = plt.plot(t[:epoch_index + 1], epoch_train_accuracy[:epoch_index + 1], label='Train accuracy')\n",
        "        l2_2, = plt.plot(t[:epoch_index + 1], epoch_test_accuracy[:epoch_index + 1], label='Test accuracy')\n",
        "        plt.legend(handles=[l2_1, l2_2])\n",
        "        plt.grid()\n",
        "        plt.gcf().gca().set_ylim(0, 100)\n",
        "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
        "#         plt.gcf().gca().set_xlim(left = 0)\n",
        "        plt.title('Accuracy')\n",
        "        display.clear_output(wait=True)\n",
        "        display.display(plt.gcf())\n",
        "\n",
        "        plt.figure(num=3)\n",
        "        plt.clf()\n",
        "        l3_1, = plt.plot(t[:epoch_index + 1], epoch_train_nochange_accuracy[:epoch_index + 1], label='Train accuracy: no change')\n",
        "        l3_2, = plt.plot(t[:epoch_index + 1], epoch_train_change_accuracy[:epoch_index + 1], label='Train accuracy: change')\n",
        "        l3_3, = plt.plot(t[:epoch_index + 1], epoch_test_nochange_accuracy[:epoch_index + 1], label='Test accuracy: no change')\n",
        "        l3_4, = plt.plot(t[:epoch_index + 1], epoch_test_change_accuracy[:epoch_index + 1], label='Test accuracy: change')\n",
        "        plt.legend(handles=[l3_1, l3_2, l3_3, l3_4])\n",
        "        plt.grid()\n",
        "        plt.gcf().gca().set_ylim(0, 100)\n",
        "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
        "#         plt.gcf().gca().set_xlim(left = 0)\n",
        "        plt.title('Accuracy per class')\n",
        "        display.clear_output(wait=True)\n",
        "        display.display(plt.gcf())\n",
        "\n",
        "        plt.figure(num=4)\n",
        "        plt.clf()\n",
        "        l4_1, = plt.plot(t[:epoch_index + 1], epoch_train_precision[:epoch_index + 1], label='Train precision')\n",
        "        l4_2, = plt.plot(t[:epoch_index + 1], epoch_train_recall[:epoch_index + 1], label='Train recall')\n",
        "        l4_3, = plt.plot(t[:epoch_index + 1], epoch_train_Fmeasure[:epoch_index + 1], label='Train Dice/F1')\n",
        "        l4_4, = plt.plot(t[:epoch_index + 1], epoch_test_precision[:epoch_index + 1], label='Test precision')\n",
        "        l4_5, = plt.plot(t[:epoch_index + 1], epoch_test_recall[:epoch_index + 1], label='Test recall')\n",
        "        l4_6, = plt.plot(t[:epoch_index + 1], epoch_test_Fmeasure[:epoch_index + 1], label='Test Dice/F1')\n",
        "        plt.legend(handles=[l4_1, l4_2, l4_3, l4_4, l4_5, l4_6])\n",
        "        plt.grid()\n",
        "        plt.gcf().gca().set_ylim(0, 1)\n",
        "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
        "#         plt.gcf().gca().set_xlim(left = 0)\n",
        "        plt.title('Precision, Recall and F-measure')\n",
        "        display.clear_output(wait=True)\n",
        "        display.display(plt.gcf())\n",
        "        \n",
        "        \n",
        "#         mean_acc = (epoch_test_nochange_accuracy[epoch_index] + epoch_test_change_accuracy[epoch_index])/2\n",
        "#         if mean_acc > best_mean_acc:\n",
        "#             best_mean_acc = mean_acc\n",
        "#             save_str = 'net-best_epoch-' + str(epoch_index + 1) + '_acc-' + str(mean_acc) + '.pth.tar'\n",
        "#             torch.save(net.state_dict(), save_str)\n",
        "        \n",
        "        \n",
        "#         fm = pr_rec[2]\n",
        "        fm = epoch_train_Fmeasure[epoch_index]\n",
        "        if fm > best_fm:\n",
        "            best_fm = fm\n",
        "            save_str = 'net-best_epoch-' + str(epoch_index + 1) + '_fm-' + str(fm) + '.pth.tar'\n",
        "            torch.save(net.state_dict(), save_str)\n",
        "        \n",
        "        lss = epoch_train_loss[epoch_index]\n",
        "        if lss < best_lss:\n",
        "            best_lss = lss\n",
        "            save_str = 'net-best_epoch-' + str(epoch_index + 1) + '_loss-' + str(lss) + '.pth.tar'\n",
        "            torch.save(net.state_dict(), save_str)\n",
        "            \n",
        "            \n",
        "#         print('Epoch loss: ' + str(tot_loss/tot_count))\n",
        "        if save:\n",
        "            im_format = 'png'\n",
        "    #         im_format = 'eps'\n",
        "\n",
        "            plt.figure(num=1)\n",
        "            plt.savefig(net_name + '-01-loss.' + im_format)\n",
        "\n",
        "            plt.figure(num=2)\n",
        "            plt.savefig(net_name + '-02-accuracy.' + im_format)\n",
        "\n",
        "            plt.figure(num=3)\n",
        "            plt.savefig(net_name + '-03-accuracy-per-class.' + im_format)\n",
        "\n",
        "            plt.figure(num=4)\n",
        "            plt.savefig(net_name + '-04-prec-rec-fmeas.' + im_format)\n",
        "        \n",
        "    out = {'train_loss': epoch_train_loss[-1],\n",
        "           'train_accuracy': epoch_train_accuracy[-1],\n",
        "           'train_nochange_accuracy': epoch_train_nochange_accuracy[-1],\n",
        "           'train_change_accuracy': epoch_train_change_accuracy[-1],\n",
        "           'test_loss': epoch_test_loss[-1],\n",
        "           'test_accuracy': epoch_test_accuracy[-1],\n",
        "           'test_nochange_accuracy': epoch_test_nochange_accuracy[-1],\n",
        "           'test_change_accuracy': epoch_test_change_accuracy[-1]}\n",
        "    \n",
        "    print('pr_c, rec_c, f_meas, pr_nc, rec_nc')\n",
        "    print(pr_rec)\n",
        "    \n",
        "    return out\n",
        "\n",
        "L = 1024\n",
        "N = 2\n",
        "\n",
        "def test(dset):\n",
        "    net.eval()\n",
        "    tot_loss = 0\n",
        "    tot_count = 0\n",
        "    tot_accurate = 0\n",
        "    \n",
        "    n = 2\n",
        "    class_correct = list(0. for i in range(n))\n",
        "    class_total = list(0. for i in range(n))\n",
        "    class_accuracy = list(0. for i in range(n))\n",
        "    \n",
        "    tp = 0\n",
        "    tn = 0\n",
        "    fp = 0\n",
        "    fn = 0\n",
        "    \n",
        "    for img_index in dset.names:\n",
        "        I1_full, I2_full, cm_full = dset.get_img(img_index)\n",
        "        \n",
        "        s = cm_full.shape\n",
        "        \n",
        "\n",
        "        steps0 = np.arange(0,s[0],ceil(s[0]/N))\n",
        "        steps1 = np.arange(0,s[1],ceil(s[1]/N))\n",
        "        for ii in range(N):\n",
        "            for jj in range(N):\n",
        "                xmin = steps0[ii]\n",
        "                if ii == N-1:\n",
        "                    xmax = s[0]\n",
        "                else:\n",
        "                    xmax = steps0[ii+1]\n",
        "                ymin = jj\n",
        "                if jj == N-1:\n",
        "                    ymax = s[1]\n",
        "                else:\n",
        "                    ymax = steps1[jj+1]\n",
        "                I1 = I1_full[:, xmin:xmax, ymin:ymax]\n",
        "                I2 = I2_full[:, xmin:xmax, ymin:ymax]\n",
        "                cm = cm_full[xmin:xmax, ymin:ymax]\n",
        "\n",
        "                I1 = Variable(torch.unsqueeze(I1, 0).float()).cuda()\n",
        "                I2 = Variable(torch.unsqueeze(I2, 0).float()).cuda()\n",
        "                cm = Variable(torch.unsqueeze(torch.from_numpy(1.0*cm),0).float()).cuda()\n",
        "\n",
        "\n",
        "                output = net(I1, I2)\n",
        "                loss = criterion(output, cm.long())\n",
        "        #         print(loss)\n",
        "                tot_loss += loss.data * np.prod(cm.size())\n",
        "                tot_count += np.prod(cm.size())\n",
        "\n",
        "                _, predicted = torch.max(output.data, 1)\n",
        "\n",
        "                c = (predicted.int() == cm.data.int())\n",
        "                for i in range(c.size(1)):\n",
        "                    for j in range(c.size(2)):\n",
        "                        l = int(cm.data[0, i, j])\n",
        "                        class_correct[l] += c[0, i, j]\n",
        "                        class_total[l] += 1\n",
        "                        \n",
        "                pr = (predicted.int() > 0).cpu().numpy()\n",
        "                gt = (cm.data.int() > 0).cpu().numpy()\n",
        "                \n",
        "                tp += np.logical_and(pr, gt).sum()\n",
        "                tn += np.logical_and(np.logical_not(pr), np.logical_not(gt)).sum()\n",
        "                fp += np.logical_and(pr, np.logical_not(gt)).sum()\n",
        "                fn += np.logical_and(np.logical_not(pr), gt).sum()\n",
        "        \n",
        "    net_loss = tot_loss/tot_count\n",
        "    net_accuracy = 100 * (tp + tn)/tot_count\n",
        "    \n",
        "    for i in range(n):\n",
        "        class_accuracy[i] = 100 * class_correct[i] / max(class_total[i],0.00001)\n",
        "\n",
        "    prec = tp / (tp + fp)\n",
        "    rec = tp / (tp + fn)\n",
        "    f_meas = 2 * prec * rec / (prec + rec)\n",
        "    prec_nc = tn / (tn + fn)\n",
        "    rec_nc = tn / (tn + fp)\n",
        "    \n",
        "    pr_rec = [prec, rec, f_meas, prec_nc, rec_nc]\n",
        "        \n",
        "    return net_loss, net_accuracy, class_accuracy, pr_rec\n",
        "    \n",
        "    \n",
        "\n",
        "\n",
        "\n",
        "    \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7q2B2ZsbtCDi",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 601
        },
        "outputId": "45266126-d631-4a63-f78e-ff465e61ca56"
      },
      "source": [
        "if LOAD_TRAINED:\n",
        "    net.load_state_dict(torch.load('net_final.pth.tar'))\n",
        "    print('LOAD OK')\n",
        "else:\n",
        "    t_start = time.time()\n",
        "    out_dic = train()\n",
        "    t_end = time.time()\n",
        "    print(out_dic)\n",
        "    print('Elapsed time:')\n",
        "    print(t_end - t_start)\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "  0%|          | 0/50 [00:00<?, ?it/s]\u001b[A\u001b[A\u001b[A\u001b[A\u001b[A\u001b[A\u001b[A\u001b[A"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch: 1 of 50\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-62-8677fed8528e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mt_start\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m     \u001b[0mout_dic\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m     \u001b[0mt_end\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout_dic\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-61-75a2db202472>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(n_epochs, save)\u001b[0m\n\u001b[1;32m     63\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 65\u001b[0;31m         \u001b[0mepoch_train_loss\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mepoch_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch_train_accuracy\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mepoch_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcl_acc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpr_rec\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_dataset\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     66\u001b[0m         \u001b[0mepoch_train_nochange_accuracy\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mepoch_index\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcl_acc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     67\u001b[0m         \u001b[0mepoch_train_change_accuracy\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mepoch_index\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcl_acc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-61-75a2db202472>\u001b[0m in \u001b[0;36mtest\u001b[0;34m(dset)\u001b[0m\n\u001b[1;32m    236\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    237\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 238\u001b[0;31m                 \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mI1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mI2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    239\u001b[0m                 \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlong\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    240\u001b[0m         \u001b[0;31m#         print(loss)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    720\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    721\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 722\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    723\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    724\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-58-7c822253e183>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x1, x2)\u001b[0m\n\u001b[1;32m     55\u001b[0m         \u001b[0mx4_0\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv4_0\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpool\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx3_0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 57\u001b[0;31m         \u001b[0mx3_1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv3_1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mx3_0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx4_0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     58\u001b[0m         \u001b[0mx2_2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv2_2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mx2_0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx3_1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     59\u001b[0m         \u001b[0mx1_3\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv1_3\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mx1_0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx2_2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: Sizes of tensors must match except in dimension 2. Got 28 and 29 (The offending index is 0)"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 0 Axes>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 0 Axes>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 0 Axes>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hpYtA84qtCDp"
      },
      "source": [
        "if not LOAD_TRAINED:\n",
        "    torch.save(net.state_dict(), 'net_final.pth.tar')\n",
        "    print('SAVE OK')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4icwmptgtCDv"
      },
      "source": [
        "\n",
        "\n",
        "def save_test_results(dset):\n",
        "    for name in tqdm(dset.names):\n",
        "        with warnings.catch_warnings():\n",
        "            I1, I2, cm = dset.get_img(name)\n",
        "            I1 = Variable(torch.unsqueeze(I1, 0).float()).cuda()\n",
        "            I2 = Variable(torch.unsqueeze(I2, 0).float()).cuda()\n",
        "            out = net(I1, I2)\n",
        "            _, predicted = torch.max(out.data, 1)\n",
        "            I = np.stack((255*cm,255*np.squeeze(predicted.cpu().numpy()),255*cm),2)\n",
        "            io.imsave(f'{net_name}-{name}.png',I)\n",
        "\n",
        "\n",
        "\n",
        "t_start = time.time()\n",
        "# save_test_results(train_dataset)\n",
        "save_test_results(test_dataset)\n",
        "t_end = time.time()\n",
        "print('Elapsed time: {}'.format(t_end - t_start))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3peRWNtxtCD0"
      },
      "source": [
        "L = 1024\n",
        "\n",
        "def kappa(tp, tn, fp, fn):\n",
        "    N = tp + tn + fp + fn\n",
        "    p0 = (tp + tn) / N\n",
        "    pe = ((tp+fp)*(tp+fn) + (tn+fp)*(tn+fn)) / (N * N)\n",
        "    \n",
        "    return (p0 - pe) / (1 - pe)\n",
        "\n",
        "def test(dset):\n",
        "    net.eval()\n",
        "    tot_loss = 0\n",
        "    tot_count = 0\n",
        "    tot_accurate = 0\n",
        "    \n",
        "    n = 2\n",
        "    class_correct = list(0. for i in range(n))\n",
        "    class_total = list(0. for i in range(n))\n",
        "    class_accuracy = list(0. for i in range(n))\n",
        "    \n",
        "    tp = 0\n",
        "    tn = 0\n",
        "    fp = 0\n",
        "    fn = 0\n",
        "    \n",
        "    for img_index in tqdm(dset.names):\n",
        "        I1_full, I2_full, cm_full = dset.get_img(img_index)\n",
        "        \n",
        "        s = cm_full.shape\n",
        "        \n",
        "        for ii in range(ceil(s[0]/L)):\n",
        "            for jj in range(ceil(s[1]/L)):\n",
        "                xmin = L*ii\n",
        "                xmax = min(L*(ii+1),s[1])\n",
        "                ymin = L*jj\n",
        "                ymax = min(L*(jj+1),s[1])\n",
        "                I1 = I1_full[:, xmin:xmax, ymin:ymax]\n",
        "                I2 = I2_full[:, xmin:xmax, ymin:ymax]\n",
        "                cm = cm_full[xmin:xmax, ymin:ymax]\n",
        "\n",
        "                I1 = Variable(torch.unsqueeze(I1, 0).float()).cuda()\n",
        "                I2 = Variable(torch.unsqueeze(I2, 0).float()).cuda()\n",
        "                cm = Variable(torch.unsqueeze(torch.from_numpy(1.0*cm),0).float()).cuda()\n",
        "\n",
        "                output = net(I1, I2)\n",
        "                    \n",
        "                loss = criterion(output, cm.long())\n",
        "                tot_loss += loss.data * np.prod(cm.size())\n",
        "                tot_count += np.prod(cm.size())\n",
        "\n",
        "                _, predicted = torch.max(output.data, 1)\n",
        "\n",
        "                c = (predicted.int() == cm.data.int())\n",
        "                for i in range(c.size(1)):\n",
        "                    for j in range(c.size(2)):\n",
        "                        l = int(cm.data[0, i, j])\n",
        "                        class_correct[l] += c[0, i, j]\n",
        "                        class_total[l] += 1\n",
        "                        \n",
        "                pr = (predicted.int() > 0).cpu().numpy()\n",
        "                gt = (cm.data.int() > 0).cpu().numpy()\n",
        "                \n",
        "                tp += np.logical_and(pr, gt).sum()\n",
        "                tn += np.logical_and(np.logical_not(pr), np.logical_not(gt)).sum()\n",
        "                fp += np.logical_and(pr, np.logical_not(gt)).sum()\n",
        "                fn += np.logical_and(np.logical_not(pr), gt).sum()\n",
        "        \n",
        "    net_loss = tot_loss/tot_count        \n",
        "    net_loss = float(net_loss.cpu().numpy())\n",
        "    \n",
        "    net_accuracy = 100 * (tp + tn)/tot_count\n",
        "    \n",
        "    for i in range(n):\n",
        "        class_accuracy[i] = 100 * class_correct[i] / max(class_total[i],0.00001)\n",
        "        class_accuracy[i] =  float(class_accuracy[i].cpu().numpy())\n",
        "\n",
        "    prec = tp / (tp + fp)\n",
        "    rec = tp / (tp + fn)\n",
        "    dice = 2 * prec * rec / (prec + rec)\n",
        "    prec_nc = tn / (tn + fn)\n",
        "    rec_nc = tn / (tn + fp)\n",
        "    \n",
        "    pr_rec = [prec, rec, dice, prec_nc, rec_nc]\n",
        "    \n",
        "    k = kappa(tp, tn, fp, fn)\n",
        "    \n",
        "    return {'net_loss': net_loss, \n",
        "            'net_accuracy': net_accuracy, \n",
        "            'class_accuracy': class_accuracy, \n",
        "            'precision': prec, \n",
        "            'recall': rec, \n",
        "            'dice': dice, \n",
        "            'kappa': k}\n",
        "\n",
        "results = test(test_dataset)\n",
        "pprint(results)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1ifeLG8otCD8"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZqSuwfg2tCED"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iZvn79GEtCER"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zHJRkQnltCEZ"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uRYCeSd9tCEg"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}