segnet-vgg16-svgnet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "# imports and stuff\n",
    "import numpy as np\n",
    "from skimage import io\n",
    "from glob import glob\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import random\n",
    "import itertools\n",
    "# Matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "# Torch imports\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data as data\n",
    "import torch.optim as optim\n",
    "import torch.optim.lr_scheduler\n",
    "import torch.nn.init\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "WINDOW_SIZE = (224, 224) # Patch size\n",
    "STRIDE = 32 # Stride for testing\n",
    "IN_CHANNELS = 5 # Number of input channels (e.g. RGB)\n",
    "FOLDER = \"../input/\" # Replace with your \"/path/to/the/ISPRS/dataset/folder/\"\n",
    "BATCH_SIZE = 15 # Number of samples in a mini-batch\n",
    "\n",
    "LABELS = [\"roads\", \"buildings\", \"low veg.\", \"trees\", \"cars\", \"clutter\"] # Label names\n",
    "N_CLASSES = len(LABELS) # Number of classes\n",
    "WEIGHTS = torch.ones(N_CLASSES) # Weights for class balancing\n",
    "CACHE = True # Store the dataset in-memory\n",
    "\n",
    "DATASET = 'Potsdam'\n",
    "\n",
    "if DATASET == 'Potsdam':\n",
    "    MAIN_FOLDER = FOLDER\n",
    "    DATA_FOLDER = MAIN_FOLDER + '3_ortho_irrg/3_Ortho_IRRG/top_potsdam_{}_IRRG.tif'\n",
    "    DSM_FOLDER = MAIN_FOLDER + '1_dsm/1_DSM/dsm_potsdam_0{}.tif'\n",
    "    NDSM_FOLDER = MAIN_FOLDER + '1_dsm_normalisation/1_DSM_normalisation/dsm_potsdam_0{}_normalized_lastools.jpg'\n",
    "    LABEL_FOLDER = MAIN_FOLDER + '5_labels_for_participants/5_Labels_for_participants/top_potsdam_{}_label.tif'\n",
    "    ERODED_FOLDER = MAIN_FOLDER + '5_labels_for_participants_no_boundary/5_Labels_for_participants_no_Boundary/top_potsdam_{}_label_noBoundary.tif'    \n",
    "elif DATASET == 'Vaihingen':\n",
    "    MAIN_FOLDER = FOLDER + 'Vaihingen/'\n",
    "    DATA_FOLDER = MAIN_FOLDER + 'top/top_mosaic_09cm_area{}.tif'\n",
    "    LABEL_FOLDER = MAIN_FOLDER + 'gts_for_participants/top_mosaic_09cm_area{}.tif'\n",
    "    ERODED_FOLDER = MAIN_FOLDER + 'gts_eroded_for_participants/top_mosaic_09cm_area{}_noBoundary.tif'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "_uuid": "a38d9cdbb61b0a0465dcf5ecc6cafeb59bfeaae9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n# We load one tile from the dataset and we display it\\nimg = io.imread(\\'../input/3_ortho_irrg/3_Ortho_IRRG/top_potsdam_2_10_IRRG.tif\\')\\nfig = plt.figure()\\nplt.subplot(2, 2, 1)\\nplt.imshow(img)\\n\\ndsm = io.imread(\\'../input/1_dsm/1_DSM/dsm_potsdam_02_10.tif\\')\\nplt.subplot(2, 2, 2)\\nplt.imshow(dsm)\\n#\\'2_12\\', \\'3_12\\', \\'5_12\\',\\'6_11\\', \\'6_12\\', \\'7_08\\'\\n#[\\'2_10\\', \\'2_11\\', \\'3_10\\', \\'3_11\\', \\'4_10\\', \\'4_11\\', \\'5_10\\', \\'6_7\\', \\'6_9\\', \\'6_10\\', \\'7_8\\', \\'7_9\\', \\'7_10\\', \\'7_11\\']\\ndsm = io.imread(\\'../input/1_dsm_normalisation/1_DSM_normalisation/dsm_potsdam_03_11_normalized_lastools.jpg\\')\\nplt.subplot(2, 2, 3)\\nplt.imshow(dsm)\\n\\n# We load the ground truth\\ngt = io.imread(\\'../input/5_labels_for_participants/5_Labels_for_participants/top_potsdam_2_10_label.tif\\')\\nplt.subplot(2, 2, 4)\\nplt.imshow(gt)\\nplt.show()\\n\\n# We also check that we can convert the ground truth into an array format\\narray_gt = convert_from_color(gt)\\nprint(\"Ground truth in numerical format has shape ({},{}) : \\n\".format(*array_gt.shape[:2]), array_gt)\\n'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ISPRS color palette\n",
    "# Let's define the standard ISPRS color palette\n",
    "\n",
    "palette = {0 : (255, 255, 255), # Impervious surfaces (white)\n",
    "           1 : (0, 0, 255),     # Buildings (blue)\n",
    "           2 : (0, 255, 255),   # Low vegetation (cyan)\n",
    "           3 : (0, 255, 0),     # Trees (green)\n",
    "           4 : (255, 255, 0),   # Cars (yellow)\n",
    "           5 : (255, 0, 0),     # Clutter (red)\n",
    "           6 : (0, 0, 0)}       # Undefined (black)\n",
    "\n",
    "invert_palette = {v: k for k, v in palette.items()}\n",
    "\n",
    "def convert_to_color(arr_2d, palette=palette):\n",
    "    \"\"\" Numeric labels to RGB-color encoding \"\"\"\n",
    "    arr_3d = np.zeros((arr_2d.shape[0], arr_2d.shape[1], 3), dtype=np.uint8)\n",
    "\n",
    "    for c, i in palette.items():\n",
    "        m = arr_2d == c\n",
    "        arr_3d[m] = i\n",
    "\n",
    "    return arr_3d\n",
    "\n",
    "def convert_from_color(arr_3d, palette=invert_palette):\n",
    "    \"\"\" RGB-color encoding to grayscale labels \"\"\"\n",
    "    arr_2d = np.zeros((arr_3d.shape[0], arr_3d.shape[1]), dtype=np.uint8)\n",
    "\n",
    "    for c, i in palette.items():\n",
    "        m = np.all(arr_3d == np.array(c).reshape(1, 1, 3), axis=2)\n",
    "        arr_2d[m] = i\n",
    "\n",
    "    return arr_2d\n",
    "'''\n",
    "# We load one tile from the dataset and we display it\n",
    "img = io.imread('../input/3_ortho_irrg/3_Ortho_IRRG/top_potsdam_2_10_IRRG.tif')\n",
    "fig = plt.figure()\n",
    "plt.subplot(2, 2, 1)\n",
    "plt.imshow(img)\n",
    "\n",
    "dsm = io.imread('../input/1_dsm/1_DSM/dsm_potsdam_02_10.tif')\n",
    "plt.subplot(2, 2, 2)\n",
    "plt.imshow(dsm)\n",
    "#'2_12', '3_12', '5_12','6_11', '6_12', '7_08'\n",
    "#['2_10', '2_11', '3_10', '3_11', '4_10', '4_11', '5_10', '6_7', '6_9', '6_10', '7_8', '7_9', '7_10', '7_11']\n",
    "dsm = io.imread('../input/1_dsm_normalisation/1_DSM_normalisation/dsm_potsdam_03_11_normalized_lastools.jpg')\n",
    "plt.subplot(2, 2, 3)\n",
    "plt.imshow(dsm)\n",
    "\n",
    "# We load the ground truth\n",
    "gt = io.imread('../input/5_labels_for_participants/5_Labels_for_participants/top_potsdam_2_10_label.tif')\n",
    "plt.subplot(2, 2, 4)\n",
    "plt.imshow(gt)\n",
    "plt.show()\n",
    "\n",
    "# We also check that we can convert the ground truth into an array format\n",
    "array_gt = convert_from_color(gt)\n",
    "print(\"Ground truth in numerical format has shape ({},{}) : \\n\".format(*array_gt.shape[:2]), array_gt)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "_uuid": "3b27f2ea65458b7cc384fa7e79bc56dba91789c2"
   },
   "outputs": [],
   "source": [
    "# Utils\n",
    "\n",
    "def get_random_pos(img, window_shape):\n",
    "    \"\"\" Extract of 2D random patch of shape window_shape in the image \"\"\"\n",
    "    w, h = window_shape\n",
    "    W, H = img.shape[-2:]\n",
    "    x1 = random.randint(0, W - w - 1)\n",
    "    x2 = x1 + w\n",
    "    y1 = random.randint(0, H - h - 1)\n",
    "    y2 = y1 + h\n",
    "    return x1, x2, y1, y2\n",
    "\n",
    "def CrossEntropy2d(input, target, weight=None, size_average=True):\n",
    "    \"\"\" 2D version of the cross entropy loss \"\"\"\n",
    "    dim = input.dim()\n",
    "    if dim == 2:\n",
    "        return F.cross_entropy(input, target, weight, size_average)\n",
    "    elif dim == 4:\n",
    "        output = input.view(input.size(0),input.size(1), -1)\n",
    "        output = torch.transpose(output,1,2).contiguous()\n",
    "        output = output.view(-1,output.size(2))\n",
    "        target = target.view(-1)\n",
    "        return F.cross_entropy(output, target,weight, size_average)\n",
    "    else:\n",
    "        raise ValueError('Expected 2 or 4 dimensions (got {})'.format(dim))\n",
    "\n",
    "def accuracy(input, target):\n",
    "    return 100 * float(np.count_nonzero(input == target)) / target.size\n",
    "\n",
    "def sliding_window(top, step=10, window_size=(20,20)):\n",
    "    \"\"\" Slide a window_shape window across the image with a stride of step \"\"\"\n",
    "    for x in range(0, top.shape[0], step):\n",
    "        if x + window_size[0] > top.shape[0]:\n",
    "            x = top.shape[0] - window_size[0]\n",
    "        for y in range(0, top.shape[1], step):\n",
    "            if y + window_size[1] > top.shape[1]:\n",
    "                y = top.shape[1] - window_size[1]\n",
    "            yield x, y, window_size[0], window_size[1]\n",
    "            \n",
    "def count_sliding_window(top, step=10, window_size=(20,20)):\n",
    "    \"\"\" Count the number of windows in an image \"\"\"\n",
    "    c = 0\n",
    "    for x in range(0, top.shape[0], step):\n",
    "        if x + window_size[0] > top.shape[0]:\n",
    "            x = top.shape[0] - window_size[0]\n",
    "        for y in range(0, top.shape[1], step):\n",
    "            if y + window_size[1] > top.shape[1]:\n",
    "                y = top.shape[1] - window_size[1]\n",
    "            c += 1\n",
    "    return c\n",
    "\n",
    "def grouper(n, iterable):\n",
    "    \"\"\" Browse an iterator by chunk of n elements \"\"\"\n",
    "    it = iter(iterable)\n",
    "    while True:\n",
    "        chunk = tuple(itertools.islice(it, n))\n",
    "        if not chunk:\n",
    "            return\n",
    "        yield chunk\n",
    "\n",
    "def metrics(predictions, gts, label_values=LABELS):\n",
    "    cm = confusion_matrix(\n",
    "            gts,\n",
    "            predictions,\n",
    "            range(len(label_values)))\n",
    "    \n",
    "    print(\"Confusion matrix :\")\n",
    "    print(cm)\n",
    "    \n",
    "    print(\"---\")\n",
    "    \n",
    "    # Compute global accuracy\n",
    "    total = sum(sum(cm))\n",
    "    accuracy = sum([cm[x][x] for x in range(len(cm))])\n",
    "    accuracy *= 100 / float(total)\n",
    "    print(\"{} pixels processed\".format(total))\n",
    "    print(\"Total accuracy : {}%\".format(accuracy))\n",
    "    \n",
    "    print(\"---\")\n",
    "    \n",
    "    # Compute F1 score\n",
    "    F1Score = np.zeros(len(label_values))\n",
    "    for i in range(len(label_values)):\n",
    "        try:\n",
    "            F1Score[i] = 2. * cm[i,i] / (np.sum(cm[i,:]) + np.sum(cm[:,i]))\n",
    "        except:\n",
    "            # Ignore exception if there is no element in class i for test set\n",
    "            pass\n",
    "    print(\"F1Score :\")\n",
    "    for l_id, score in enumerate(F1Score):\n",
    "        print(\"{}: {}\".format(label_values[l_id], score))\n",
    "\n",
    "    print(\"---\")\n",
    "        \n",
    "    # Compute kappa coefficient\n",
    "    total = np.sum(cm)\n",
    "    pa = np.trace(cm) / float(total)\n",
    "    pe = np.sum(np.sum(cm, axis=0) * np.sum(cm, axis=1)) / float(total*total)\n",
    "    kappa = (pa - pe) / (1 - pe);\n",
    "    print(\"Kappa: \" + str(kappa))\n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "_uuid": "806064dbf5367ecca8b42065a42e3a36285ccc56"
   },
   "outputs": [],
   "source": [
    "# Dataset class\n",
    "\n",
    "class ISPRS_dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, train_ids, dsm_ids, data_files=DATA_FOLDER, label_files=LABEL_FOLDER,\n",
    "                dsm_files=DSM_FOLDER, ndsm_files=NDSM_FOLDER,\n",
    "                cache=False, augmentation=True):\n",
    "        super(ISPRS_dataset, self).__init__()\n",
    "        \n",
    "        self.augmentation = augmentation\n",
    "        self.cache = cache\n",
    "        \n",
    "        # List of files\n",
    "        self.data_files = [DATA_FOLDER.format(id) for id in train_ids]\n",
    "        self.dsm_files = [DSM_FOLDER.format(id) for id in dsm_ids]\n",
    "        self.ndsm_files = [NDSM_FOLDER.format(id) for id in dsm_ids]\n",
    "        self.label_files = [LABEL_FOLDER.format(id) for id in train_ids]\n",
    "        \n",
    "        '''\n",
    "        # Sanity check : raise an error if some files do not exist\n",
    "        for f in self.data_files + self.label_files + self.dsm_files + self.ndsm_files:\n",
    "            if not os.path.isfile(f):\n",
    "                f1 = f[:36] + '0' + f[36:]\n",
    "                if not os.path.isfile(f1):\n",
    "                    f2 = f[:64] + '0' + f[64:]\n",
    "                    if not os.path.isfile(f2):\n",
    "                        raise KeyError('{} is not a file !'.format(f2))\n",
    "        '''\n",
    "        \n",
    "        # Initialize cache dicts\n",
    "        self.data_cache_ = {}\n",
    "        self.label_cache_ = {}\n",
    "            \n",
    "    \n",
    "    def __len__(self):\n",
    "        # Default epoch size is 10000 samples\n",
    "        return 5000\n",
    "    \n",
    "    @classmethod\n",
    "    def data_augmentation(cls, *arrays, flip=True, mirror=True):\n",
    "        will_flip, will_mirror = False, False\n",
    "        if flip and random.random() < 0.5:\n",
    "            will_flip = True\n",
    "        if mirror and random.random() < 0.5:\n",
    "            will_mirror = True\n",
    "        \n",
    "        results = []\n",
    "        for array in arrays:\n",
    "            if will_flip:\n",
    "                if len(array.shape) == 2:\n",
    "                    array = array[::-1, :]\n",
    "                else:\n",
    "                    array = array[:, ::-1, :]\n",
    "            if will_mirror:\n",
    "                if len(array.shape) == 2:\n",
    "                    array = array[:, ::-1]\n",
    "                else:\n",
    "                    array = array[:, :, ::-1]\n",
    "            results.append(np.copy(array))\n",
    "            \n",
    "        return tuple(results)\n",
    "    \n",
    "    def __getitem__(self, i):\n",
    "        # Pick a random image\n",
    "        random_idx = random.randint(0, len(self.data_files) - 1)\n",
    "        \n",
    "        # If the tile hasn't been loaded yet, put in cache\n",
    "        if random_idx in self.data_cache_.keys():\n",
    "            data = self.data_cache_[random_idx]\n",
    "        else:\n",
    "            # Data is normalized in [0, 1]\n",
    "            im = np.dstack((io.imread(self.data_files[random_idx]), io.imread(self.dsm_files[random_idx])))\n",
    "            im = np.dstack((im, io.imread(self.ndsm_files[random_idx])))\n",
    "            data = 1/255 * np.asarray(im.transpose((2,0,1)), dtype='float32')\n",
    "            if self.cache:\n",
    "                self.data_cache_[random_idx] = data\n",
    "            \n",
    "        if random_idx in self.label_cache_.keys():\n",
    "            label = self.label_cache_[random_idx]\n",
    "        else: \n",
    "            # Labels are converted from RGB to their numeric values\n",
    "            label = np.asarray(convert_from_color(io.imread(self.label_files[random_idx])), dtype='int64')\n",
    "            if self.cache:\n",
    "                self.label_cache_[random_idx] = label\n",
    "\n",
    "        # Get a random patch\n",
    "        x1, x2, y1, y2 = get_random_pos(data, WINDOW_SIZE)\n",
    "        data_p = data[:, x1:x2,y1:y2]\n",
    "        label_p = label[x1:x2,y1:y2]\n",
    "        \n",
    "        # Data augmentation\n",
    "        data_p, label_p = self.data_augmentation(data_p, label_p)\n",
    "\n",
    "        # Return the torch.Tensor values\n",
    "        return (torch.from_numpy(data_p),\n",
    "                torch.from_numpy(label_p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "_uuid": "a4780be561a9f6912bcbede2051ef58fb9f31c13"
   },
   "outputs": [],
   "source": [
    "class SegNet(nn.Module):\n",
    "    # SegNet network\n",
    "    @staticmethod\n",
    "    def weight_init(m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            torch.nn.init.kaiming_normal(m.weight.data)\n",
    "    \n",
    "    def __init__(self, in_channels=IN_CHANNELS, out_channels=N_CLASSES):\n",
    "        super(SegNet, self).__init__()\n",
    "        self.pool = nn.MaxPool2d(2, return_indices=True)\n",
    "        self.unpool = nn.MaxUnpool2d(2)\n",
    "        \n",
    "        self.conv1_1 = nn.Conv2d(in_channels, 64, 3, padding=1)\n",
    "        self.conv1_1_bn = nn.BatchNorm2d(64)\n",
    "        self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1)\n",
    "        self.conv1_2_bn = nn.BatchNorm2d(64)\n",
    "        \n",
    "        self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1)\n",
    "        self.conv2_1_bn = nn.BatchNorm2d(128)\n",
    "        self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1)\n",
    "        self.conv2_2_bn = nn.BatchNorm2d(128)\n",
    "        \n",
    "        self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1)\n",
    "        self.conv3_1_bn = nn.BatchNorm2d(256)\n",
    "        self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.conv3_2_bn = nn.BatchNorm2d(256)\n",
    "        self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.conv3_3_bn = nn.BatchNorm2d(256)\n",
    "        \n",
    "        self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1)\n",
    "        self.conv4_1_bn = nn.BatchNorm2d(512)\n",
    "        self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv4_2_bn = nn.BatchNorm2d(512)\n",
    "        self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv4_3_bn = nn.BatchNorm2d(512)\n",
    "        \n",
    "        self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv5_1_bn = nn.BatchNorm2d(512)\n",
    "        self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv5_2_bn = nn.BatchNorm2d(512)\n",
    "        self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv5_3_bn = nn.BatchNorm2d(512)\n",
    "        \n",
    "        self.conv5_3_D = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv5_3_D_bn = nn.BatchNorm2d(512)\n",
    "        self.conv5_2_D = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv5_2_D_bn = nn.BatchNorm2d(512)\n",
    "        self.conv5_1_D = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv5_1_D_bn = nn.BatchNorm2d(512)\n",
    "        \n",
    "        self.conv4_3_D = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv4_3_D_bn = nn.BatchNorm2d(512)\n",
    "        self.conv4_2_D = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv4_2_D_bn = nn.BatchNorm2d(512)\n",
    "        self.conv4_1_D = nn.Conv2d(512, 256, 3, padding=1)\n",
    "        self.conv4_1_D_bn = nn.BatchNorm2d(256)\n",
    "        \n",
    "        self.conv3_3_D = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.conv3_3_D_bn = nn.BatchNorm2d(256)\n",
    "        self.conv3_2_D = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.conv3_2_D_bn = nn.BatchNorm2d(256)\n",
    "        self.conv3_1_D = nn.Conv2d(256, 128, 3, padding=1)\n",
    "        self.conv3_1_D_bn = nn.BatchNorm2d(128)\n",
    "        \n",
    "        self.conv2_2_D = nn.Conv2d(128, 128, 3, padding=1)\n",
    "        self.conv2_2_D_bn = nn.BatchNorm2d(128)\n",
    "        self.conv2_1_D = nn.Conv2d(128, 64, 3, padding=1)\n",
    "        self.conv2_1_D_bn = nn.BatchNorm2d(64)\n",
    "        \n",
    "        self.conv1_2_D = nn.Conv2d(64, 64, 3, padding=1)\n",
    "        self.conv1_2_D_bn = nn.BatchNorm2d(64)\n",
    "        self.conv1_1_D = nn.Conv2d(64, out_channels, 3, padding=1)\n",
    "        \n",
    "        self.apply(self.weight_init)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # Encoder block 1\n",
    "        x = self.conv1_1_bn(F.relu(self.conv1_1(x)))\n",
    "        x = self.conv1_2_bn(F.relu(self.conv1_2(x)))\n",
    "        x, mask1 = self.pool(x)\n",
    "        \n",
    "        # Encoder block 2\n",
    "        x = self.conv2_1_bn(F.relu(self.conv2_1(x)))\n",
    "        x = self.conv2_2_bn(F.relu(self.conv2_2(x)))\n",
    "        x, mask2 = self.pool(x)\n",
    "        \n",
    "        # Encoder block 3\n",
    "        x = self.conv3_1_bn(F.relu(self.conv3_1(x)))\n",
    "        x = self.conv3_2_bn(F.relu(self.conv3_2(x)))\n",
    "        x = self.conv3_3_bn(F.relu(self.conv3_3(x)))\n",
    "        x, mask3 = self.pool(x)\n",
    "        \n",
    "        # Encoder block 4\n",
    "        x = self.conv4_1_bn(F.relu(self.conv4_1(x)))\n",
    "        x = self.conv4_2_bn(F.relu(self.conv4_2(x)))\n",
    "        x = self.conv4_3_bn(F.relu(self.conv4_3(x)))\n",
    "        x, mask4 = self.pool(x)\n",
    "        \n",
    "        # Encoder block 5\n",
    "        x = self.conv5_1_bn(F.relu(self.conv5_1(x)))\n",
    "        x = self.conv5_2_bn(F.relu(self.conv5_2(x)))\n",
    "        x = self.conv5_3_bn(F.relu(self.conv5_3(x)))\n",
    "        x, mask5 = self.pool(x)\n",
    "        \n",
    "        # Decoder block 5\n",
    "        x = self.unpool(x, mask5)\n",
    "        x = self.conv5_3_D_bn(F.relu(self.conv5_3_D(x)))\n",
    "        x = self.conv5_2_D_bn(F.relu(self.conv5_2_D(x)))\n",
    "        x = self.conv5_1_D_bn(F.relu(self.conv5_1_D(x)))\n",
    "        \n",
    "        # Decoder block 4\n",
    "        x = self.unpool(x, mask4)\n",
    "        x = self.conv4_3_D_bn(F.relu(self.conv4_3_D(x)))\n",
    "        x = self.conv4_2_D_bn(F.relu(self.conv4_2_D(x)))\n",
    "        x = self.conv4_1_D_bn(F.relu(self.conv4_1_D(x)))\n",
    "        \n",
    "        # Decoder block 3\n",
    "        x = self.unpool(x, mask3)\n",
    "        x = self.conv3_3_D_bn(F.relu(self.conv3_3_D(x)))\n",
    "        x = self.conv3_2_D_bn(F.relu(self.conv3_2_D(x)))\n",
    "        x = self.conv3_1_D_bn(F.relu(self.conv3_1_D(x)))\n",
    "        \n",
    "        # Decoder block 2\n",
    "        x = self.unpool(x, mask2)\n",
    "        x = self.conv2_2_D_bn(F.relu(self.conv2_2_D(x)))\n",
    "        x = self.conv2_1_D_bn(F.relu(self.conv2_1_D(x)))\n",
    "        \n",
    "        # Decoder block 1\n",
    "        x = self.unpool(x, mask1)\n",
    "        x = self.conv1_2_D_bn(F.relu(self.conv1_2_D(x)))\n",
    "        x = F.log_softmax(self.conv1_1_D(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "_uuid": "8e48c6afec943195913b292c4786a9adf01e584b"
   },
   "outputs": [],
   "source": [
    "# instantiate the network\n",
    "net = SegNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "_uuid": "13a73f4984480194d6ea660ae881b84d93fbb0e1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping features.0.weight to conv1_1.weight\n",
      "Mapping features.0.bias to conv1_1.bias\n",
      "Mapping features.1.weight to conv1_1_bn.weight\n",
      "Mapping features.1.bias to conv1_1_bn.bias\n",
      "Mapping features.1.running_mean to conv1_1_bn.running_mean\n",
      "Mapping features.1.running_var to conv1_1_bn.running_var\n",
      "Mapping features.3.weight to conv1_1_bn.num_batches_tracked\n",
      "Mapping features.3.bias to conv1_2.weight\n",
      "Mapping features.4.weight to conv1_2.bias\n",
      "Mapping features.4.bias to conv1_2_bn.weight\n",
      "Mapping features.4.running_mean to conv1_2_bn.bias\n",
      "Mapping features.4.running_var to conv1_2_bn.running_mean\n",
      "Mapping features.7.weight to conv1_2_bn.running_var\n",
      "Mapping features.7.bias to conv1_2_bn.num_batches_tracked\n",
      "Mapping features.8.weight to conv2_1.weight\n",
      "Mapping features.8.bias to conv2_1.bias\n",
      "Mapping features.8.running_mean to conv2_1_bn.weight\n",
      "Mapping features.8.running_var to conv2_1_bn.bias\n",
      "Mapping features.10.weight to conv2_1_bn.running_mean\n",
      "Mapping features.10.bias to conv2_1_bn.running_var\n",
      "Mapping features.11.weight to conv2_1_bn.num_batches_tracked\n",
      "Mapping features.11.bias to conv2_2.weight\n",
      "Mapping features.11.running_mean to conv2_2.bias\n",
      "Mapping features.11.running_var to conv2_2_bn.weight\n",
      "Mapping features.14.weight to conv2_2_bn.bias\n",
      "Mapping features.14.bias to conv2_2_bn.running_mean\n",
      "Mapping features.15.weight to conv2_2_bn.running_var\n",
      "Mapping features.15.bias to conv2_2_bn.num_batches_tracked\n",
      "Mapping features.15.running_mean to conv3_1.weight\n",
      "Mapping features.15.running_var to conv3_1.bias\n",
      "Mapping features.17.weight to conv3_1_bn.weight\n",
      "Mapping features.17.bias to conv3_1_bn.bias\n",
      "Mapping features.18.weight to conv3_1_bn.running_mean\n",
      "Mapping features.18.bias to conv3_1_bn.running_var\n",
      "Mapping features.18.running_mean to conv3_1_bn.num_batches_tracked\n",
      "Mapping features.18.running_var to conv3_2.weight\n",
      "Mapping features.20.weight to conv3_2.bias\n",
      "Mapping features.20.bias to conv3_2_bn.weight\n",
      "Mapping features.21.weight to conv3_2_bn.bias\n",
      "Mapping features.21.bias to conv3_2_bn.running_mean\n",
      "Mapping features.21.running_mean to conv3_2_bn.running_var\n",
      "Mapping features.21.running_var to conv3_2_bn.num_batches_tracked\n",
      "Mapping features.24.weight to conv3_3.weight\n",
      "Mapping features.24.bias to conv3_3.bias\n",
      "Mapping features.25.weight to conv3_3_bn.weight\n",
      "Mapping features.25.bias to conv3_3_bn.bias\n",
      "Mapping features.25.running_mean to conv3_3_bn.running_mean\n",
      "Mapping features.25.running_var to conv3_3_bn.running_var\n",
      "Mapping features.27.weight to conv3_3_bn.num_batches_tracked\n",
      "Mapping features.27.bias to conv4_1.weight\n",
      "Mapping features.28.weight to conv4_1.bias\n",
      "Mapping features.28.bias to conv4_1_bn.weight\n",
      "Mapping features.28.running_mean to conv4_1_bn.bias\n",
      "Mapping features.28.running_var to conv4_1_bn.running_mean\n",
      "Mapping features.30.weight to conv4_1_bn.running_var\n",
      "Mapping features.30.bias to conv4_1_bn.num_batches_tracked\n",
      "Mapping features.31.weight to conv4_2.weight\n",
      "Mapping features.31.bias to conv4_2.bias\n",
      "Mapping features.31.running_mean to conv4_2_bn.weight\n",
      "Mapping features.31.running_var to conv4_2_bn.bias\n",
      "Mapping features.34.weight to conv4_2_bn.running_mean\n",
      "Mapping features.34.bias to conv4_2_bn.running_var\n",
      "Mapping features.35.weight to conv4_2_bn.num_batches_tracked\n",
      "Mapping features.35.bias to conv4_3.weight\n",
      "Mapping features.35.running_mean to conv4_3.bias\n",
      "Mapping features.35.running_var to conv4_3_bn.weight\n",
      "Mapping features.37.weight to conv4_3_bn.bias\n",
      "Mapping features.37.bias to conv4_3_bn.running_mean\n",
      "Mapping features.38.weight to conv4_3_bn.running_var\n",
      "Mapping features.38.bias to conv4_3_bn.num_batches_tracked\n",
      "Mapping features.38.running_mean to conv5_1.weight\n",
      "Mapping features.38.running_var to conv5_1.bias\n",
      "Mapping features.40.weight to conv5_1_bn.weight\n",
      "Mapping features.40.bias to conv5_1_bn.bias\n",
      "Mapping features.41.weight to conv5_1_bn.running_mean\n",
      "Mapping features.41.bias to conv5_1_bn.running_var\n",
      "Mapping features.41.running_mean to conv5_1_bn.num_batches_tracked\n",
      "Mapping features.41.running_var to conv5_2.weight\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "try:\n",
    "    from urllib.request import URLopener\n",
    "except ImportError:\n",
    "    from urllib import URLopener\n",
    "\n",
    "# Download VGG-16 weights from PyTorch\n",
    "vgg_url = 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth'\n",
    "if not os.path.isfile('./vgg16_bn-6c64b313.pth'):\n",
    "    weights = URLopener().retrieve(vgg_url, './vgg16_bn-6c64b313.pth')\n",
    "\n",
    "vgg16_weights = torch.load('./vgg16_bn-6c64b313.pth')\n",
    "mapped_weights = {}\n",
    "for k_vgg, k_segnet in zip(vgg16_weights.keys(), net.state_dict().keys()):\n",
    "    if \"features\" in k_vgg:\n",
    "        mapped_weights[k_segnet] = vgg16_weights[k_vgg]\n",
    "        print(\"Mapping {} to {}\".format(k_vgg, k_segnet))\n",
    "        \n",
    "try:\n",
    "    net.load_state_dict(mapped_weights)\n",
    "    print(\"Loaded VGG-16 weights in SegNet !\")\n",
    "except:\n",
    "    # Ignore missing keys\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "_uuid": "b95ea537bb3c483b4b590375207cc373ed65f185"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SegNet(\n",
       "  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  (unpool): MaxUnpool2d(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))\n",
       "  (conv1_1): Conv2d(5, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv1_1_bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv1_2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv1_2_bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv2_1): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv2_1_bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv2_2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv2_2_bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3_1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv3_1_bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3_2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv3_2_bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3_3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv3_3_bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv4_1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv4_1_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv4_2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv4_2_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv4_3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv4_3_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5_1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv5_1_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5_2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv5_2_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5_3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv5_3_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5_3_D): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv5_3_D_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5_2_D): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv5_2_D_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5_1_D): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv5_1_D_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      