{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3AbTeDP5Tbou"
      },
      "source": [
        "# Segmentation\n",
        "\n",
        "We have already learnt about Object Detection, which allows us to locate objects in the image by predicting their *bounding boxes*. However, for some tasks we do not only need bounding boxes, but also more precise object localization. This task is called  **segmentation**.\n",
        "\n",
        "Segmentation can be viewed as **pixel classification**, whereas for **each** pixel of image we must predict its class (*background* being one of the classes). There are two main segmentation algorithms:\n",
        "\n",
        "* **Semantic segmentation** only tells pixel class, and does not make a distinction between different objects of the same class\n",
        "* **Instance segmentation** divides classes into different instances. \n",
        "\n",
        "For instance segmentation 10 sheep are different objects, for semantic segmentation all sheep are represented by one class.\n",
        "\n",
        "<img src=\"images/instance_vs_semantic.jpeg\" width=\"50%\">\n",
        "\n",
        "> Image from [this blog post](https://nirmalamurali.medium.com/image-classification-vs-semantic-segmentation-vs-instance-segmentation-625c33a08d50)\n",
        "\n",
        "There are different neural architectures for segmentation, but they all have the same structure:\n",
        "\n",
        "* **Encoder** extracts features from input image\n",
        "* **Decoder** transforms those features into the **mask image**, with the same size and number of channels corresponding to the number of classes.\n",
        "\n",
        "<img src=\"images/segm.png\" width=\"80%\">\n",
        "\n",
        "> Image from [this publication](https://arxiv.org/pdf/2001.05566.pdf)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Prerequsites\n",
        "\n",
        "To begin with, we will import required libraries, and check if there is GPU available for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T15:48:06.861688Z",
          "iopub.status.busy": "2022-04-08T15:48:06.861254Z",
          "iopub.status.idle": "2022-04-08T15:48:06.868219Z",
          "shell.execute_reply": "2022-04-08T15:48:06.867567Z",
          "shell.execute_reply.started": "2022-04-08T15:48:06.861644Z"
        },
        "id": "tv1T3XemFMpE",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import matplotlib.pyplot as plt\n",
        "from torchvision import transforms\n",
        "from torch import nn\n",
        "from torch import optim\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "from skimage.io import imread\n",
        "from skimage.transform import resize\n",
        "import os\n",
        "torch.manual_seed(42)\n",
        "np.random.seed(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:03:32.933989Z",
          "iopub.status.busy": "2022-04-08T16:03:32.933733Z",
          "iopub.status.idle": "2022-04-08T16:03:32.937996Z",
          "shell.execute_reply": "2022-04-08T16:03:32.937366Z",
          "shell.execute_reply.started": "2022-04-08T16:03:32.933959Z"
        },
        "id": "xhSEogpDFMpK",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "train_size = 0.9\n",
        "lr = 1e-3\n",
        "weight_decay = 1e-6\n",
        "batch_size = 32\n",
        "epochs = 30"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D4if75qwFMpJ"
      },
      "source": [
        "## The Dataset\n",
        "\n",
        "We will use the <a href=\"https://www.fc.up.pt/addi/ph2%20database.html\">PH<sup>2</sup> Database</a> of dermoscopy images of human nevi. This dataset contains 200 images of three classes: typical nevus, atypical nevus, and melanoma. All images also contain corresponding **mask** that outline the nevus.\n",
        "\n",
        "The code below downloads the dataset from the original location and decompresses it. You would need to have `unrar` utility installed in order for this code to work, you may install it using `sudo apt-get install unrar` on Linux, or by downloading command-line version for Windows [here](https://www.rarlab.com/rar_add.htm)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:47:25.071036Z",
          "iopub.status.busy": "2022-04-08T16:47:25.070746Z",
          "iopub.status.idle": "2022-04-08T16:47:32.612115Z",
          "shell.execute_reply": "2022-04-08T16:47:32.611253Z",
          "shell.execute_reply.started": "2022-04-08T16:47:25.071005Z"
        },
        "id": "TkuYGLRKFMpK",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "#!apt-get install rar\n",
        "!wget https://www.dropbox.com/s/k88qukc20ljnbuo/PH2Dataset.rar\n",
        "!unrar x -Y PH2Dataset.rar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we will define the code to load the dataset. We will transform all images into 256x256 size, and split the dataset into train and test part. This function returns train and test datasets, each containing original images and masks outlining the nevus."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "Rumy9ldAFteW"
      },
      "outputs": [],
      "source": [
        "def load_dataset(train_part, root='PH2Dataset'):\n",
        "    images = []\n",
        "    masks = []\n",
        "\n",
        "    for root, dirs, files in os.walk(os.path.join(root, 'PH2 Dataset images')):\n",
        "        if root.endswith('_Dermoscopic_Image'):\n",
        "            images.append(imread(os.path.join(root, files[0])))\n",
        "        if root.endswith('_lesion'):\n",
        "            masks.append(imread(os.path.join(root, files[0])))\n",
        "\n",
        "    size = (256, 256)\n",
        "    images = torch.permute(torch.FloatTensor(np.array([resize(image, size, mode='constant', anti_aliasing=True,) for image in images])), (0, 3, 1, 2))\n",
        "    masks = torch.FloatTensor(np.array([resize(mask, size, mode='constant', anti_aliasing=False) > 0.5 for mask in masks])).unsqueeze(1)\n",
        "\n",
        "    indices = np.random.permutation(range(len(images)))\n",
        "    train_part = int(train_part * len(images))\n",
        "    train_ind = indices[:train_part]\n",
        "    test_ind = indices[train_part:]\n",
        "\n",
        "    train_dataset = (images[train_ind, :, :, :], masks[train_ind, :, :, :])\n",
        "    test_dataset = (images[test_ind, :, :, :], masks[test_ind, :, :, :])\n",
        "\n",
        "    return train_dataset, test_dataset\n",
        "\n",
        "train_dataset, test_dataset = load_dataset(train_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's now plot some of the images from the dataset to see how they look like:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:03:58.259127Z",
          "iopub.status.busy": "2022-04-08T16:03:58.258819Z",
          "iopub.status.idle": "2022-04-08T16:03:58.266433Z",
          "shell.execute_reply": "2022-04-08T16:03:58.265489Z",
          "shell.execute_reply.started": "2022-04-08T16:03:58.259090Z"
        },
        "id": "jP2_-AjIFMpO",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "def plotn(n, data, only_mask=False):\n",
        "    images, masks = data[0], data[1]\n",
        "    fig, ax = plt.subplots(1, n)\n",
        "    fig1, ax1 = plt.subplots(1, n)\n",
        "    for i, (img, mask) in enumerate(zip(images, masks)):\n",
        "        if i == n:\n",
        "            break\n",
        "        if not only_mask:\n",
        "            ax[i].imshow(torch.permute(img, (1, 2, 0)))\n",
        "        else:\n",
        "            ax[i].imshow(img[0])\n",
        "        ax1[i].imshow(mask[0])\n",
        "        ax[i].axis('off')\n",
        "        ax1[i].axis('off')\n",
        "    plt.show()\n",
        "\n",
        "plotn(5, train_dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will also need dataloaders to feed the data into our neural network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:03:58.247264Z",
          "iopub.status.busy": "2022-04-08T16:03:58.246853Z",
          "iopub.status.idle": "2022-04-08T16:03:58.256661Z",
          "shell.execute_reply": "2022-04-08T16:03:58.255964Z",
          "shell.execute_reply.started": "2022-04-08T16:03:58.247222Z"
        },
        "id": "rUGDAa61FMpN",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "train_dataloader = torch.utils.data.DataLoader(list(zip(train_dataset[0], train_dataset[1])), batch_size=batch_size, shuffle=True)\n",
        "test_dataloader = torch.utils.data.DataLoader(list(zip(test_dataset[0], test_dataset[1])), batch_size=1, shuffle=False)\n",
        "dataloaders = (train_dataloader, test_dataloader)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ORmas8XhYfS8"
      },
      "source": [
        "## SegNet\n",
        "\n",
        "The simplest encoder-decoder architecture is called **SegNet**. It uses standard CNN with convolutions and poolings in the encoder, and deconvolution CNN that includes convolutions and upsamplings in decoder. It also relies on batch normalization to train multi-layered network successfully.\n",
        "\n",
        "<img src=\"images/segnet.png\" width=\"80%\">\n",
        "\n",
        "> Image from this paper: Badrinarayanan, V., Kendall, A., & Cipolla, R. (2015). [SegNet: A deep convolutional\n",
        "encoder-decoder architecture for image segmentation](https://arxiv.org/pdf/1511.00561.pdf)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:03:59.239435Z",
          "iopub.status.busy": "2022-04-08T16:03:59.239186Z",
          "iopub.status.idle": "2022-04-08T16:03:59.257319Z",
          "shell.execute_reply": "2022-04-08T16:03:59.256388Z",
          "shell.execute_reply.started": "2022-04-08T16:03:59.239401Z"
        },
        "id": "QDsSrmbeTbp9",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class SegNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.enc_conv0 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=(3,3), padding=1)\n",
        "        self.act0 = nn.ReLU()\n",
        "        self.bn0 = nn.BatchNorm2d(16)\n",
        "        self.pool0 = nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.enc_conv1 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3,3), padding=1)\n",
        "        self.act1 = nn.ReLU()\n",
        "        self.bn1 = nn.BatchNorm2d(32)\n",
        "        self.pool1 = nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.enc_conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3,3), padding=1)\n",
        "        self.act2 = nn.ReLU()\n",
        "        self.bn2 = nn.BatchNorm2d(64)\n",
        "        self.pool2 =  nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.enc_conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3,3), padding=1)\n",
        "        self.act3 = nn.ReLU()\n",
        "        self.bn3 = nn.BatchNorm2d(128)\n",
        "        self.pool3 =  nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.bottleneck_conv = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3,3), padding=1)\n",
        "        \n",
        "        self.upsample0 =  nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        self.dec_conv0 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=(3,3), padding=1)\n",
        "        self.dec_act0 = nn.ReLU()\n",
        "        self.dec_bn0 = nn.BatchNorm2d(128)\n",
        "\n",
        "        self.upsample1 =  nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        self.dec_conv1 =  nn.Conv2d(in_channels=128, out_channels=64, kernel_size=(3,3), padding=1)\n",
        "        self.dec_act1 = nn.ReLU()\n",
        "        self.dec_bn1 = nn.BatchNorm2d(64)\n",
        "\n",
        "        self.upsample2 = nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        \n",
        "        self.dec_conv2 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=(3,3), padding=1)\n",
        "        self.dec_act2 = nn.ReLU()\n",
        "        self.dec_bn2 = nn.BatchNorm2d(32)\n",
        "\n",
        "        self.upsample3 = nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        self.dec_conv3 = nn.Conv2d(in_channels=32, out_channels=1, kernel_size=(1,1))\n",
        "\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "\n",
        "    def forward(self, x):\n",
        "        e0 = self.pool0(self.bn0(self.act0(self.enc_conv0(x))))\n",
        "        e1 = self.pool1(self.bn1(self.act1(self.enc_conv1(e0))))\n",
        "        e2 = self.pool2(self.bn2(self.act2(self.enc_conv2(e1))))\n",
        "        e3 = self.pool3(self.bn3(self.act3(self.enc_conv3(e2))))\n",
        "\n",
        "        b = self.bottleneck_conv(e3)\n",
        "\n",
        "        d0 = self.dec_bn0(self.dec_act0(self.dec_conv0(self.upsample0(b))))\n",
        "        d1 = self.dec_bn1(self.dec_act1(self.dec_conv1(self.upsample1(d0))))\n",
        "        d2 = self.dec_bn2(self.dec_act2(self.dec_conv2(self.upsample2(d1))))\n",
        "        d3 = self.sigmoid(self.dec_conv3(self.upsample3(d2)))\n",
        "        return d3"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We should especially mention the loss function that is used for segmentation. In classical autoencoders we need to measure the similarity between two images, and we can use mean square error to do that. In segmentation, each pixel in the target mask image represents the class number (one-hot-encoded along the third dimension), so we need to use loss functions specific for classification - cross-entropy loss, averaged over all pixels. If the mask is binary (as in our example) - we will use **binary cross-entropy loss** (BCE).   "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:03:59.259154Z",
          "iopub.status.busy": "2022-04-08T16:03:59.258856Z",
          "iopub.status.idle": "2022-04-08T16:03:59.284201Z",
          "shell.execute_reply": "2022-04-08T16:03:59.283559Z",
          "shell.execute_reply.started": "2022-04-08T16:03:59.259108Z"
        },
        "id": "2GoR8-huFMpn",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "model = SegNet().to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
        "loss_fn = nn.BCEWithLogitsLoss()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Training loop is defined in the usual way:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:03:59.286008Z",
          "iopub.status.busy": "2022-04-08T16:03:59.285826Z",
          "iopub.status.idle": "2022-04-08T16:03:59.298404Z",
          "shell.execute_reply": "2022-04-08T16:03:59.297656Z",
          "shell.execute_reply.started": "2022-04-08T16:03:59.285986Z"
        },
        "id": "HQ_UFglyTbqM",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "def train(dataloaders, model, loss_fn, optimizer, epochs, device):\n",
        "    tqdm_iter = tqdm(range(epochs))\n",
        "    train_dataloader, test_dataloader = dataloaders[0], dataloaders[1]\n",
        "\n",
        "    for epoch in tqdm_iter:\n",
        "        model.train()\n",
        "        train_loss = 0.0\n",
        "        test_loss = 0.0\n",
        "\n",
        "        for batch in train_dataloader:\n",
        "            imgs, labels = batch\n",
        "            imgs = imgs.to(device)\n",
        "            labels = labels.to(device)\n",
        "\n",
        "            preds = model(imgs)\n",
        "            loss = loss_fn(preds, labels)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            train_loss += loss.item()\n",
        "\n",
        "        model.eval()\n",
        "        with torch.no_grad():\n",
        "            for batch in test_dataloader:\n",
        "                imgs, labels = batch\n",
        "                imgs = imgs.to(device)\n",
        "                labels = labels.to(device)\n",
        "\n",
        "                preds = model(imgs)\n",
        "                loss = loss_fn(preds, labels)\n",
        "\n",
        "                test_loss += loss.item()\n",
        "\n",
        "        train_loss /= len(train_dataloader)\n",
        "        test_loss /= len(test_dataloader)\n",
        "\n",
        "        tqdm_dct = {'train loss:': train_loss, 'test loss:': test_loss}\n",
        "        tqdm_iter.set_postfix(tqdm_dct, refresh=True)\n",
        "        tqdm_iter.refresh()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "execution": {
          "iopub.execute_input": "2022-04-08T16:03:59.302207Z",
          "iopub.status.busy": "2022-04-08T16:03:59.301979Z",
          "iopub.status.idle": "2022-04-08T16:17:44.792683Z",
          "shell.execute_reply": "2022-04-08T16:17:44.791975Z",
          "shell.execute_reply.started": "2022-04-08T16:03:59.302184Z"
        },
        "id": "MsgM_kZRFMpo",
        "outputId": "8ff2de4a-ad8d-4b57-bdeb-ecaa90f742e2",
        "trusted": true
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 30/30 [16:01<00:00, 32.04s/it, train loss:=0.593, test loss:=0.577]\n"
          ]
        }
      ],
      "source": [
        "train(dataloaders, model, loss_fn, optimizer, epochs, device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To evaluate our model, we will just plot target masks and predicted masks for a number of images:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 203
        },
        "execution": {
          "iopub.execute_input": "2022-04-08T16:17:44.795590Z",
          "iopub.status.busy": "2022-04-08T16:17:44.794942Z",
          "iopub.status.idle": "2022-04-08T16:17:45.663916Z",
          "shell.execute_reply": "2022-04-08T16:17:45.663160Z",
          "shell.execute_reply.started": "2022-04-08T16:17:44.795550Z"
        },
        "id": "FXvR7P1FFMpo",
        "outputId": "16b1f56d-93e0-48a0-d0c2-a8214b537741",
        "trusted": true
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 7948 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 7048 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "model.eval()\n",
        "predictions = []\n",
        "image_mask = []\n",
        "plots = 5\n",
        "images, masks = test_dataset[0], test_dataset[1]\n",
        "for i, (img, mask) in enumerate(zip(images, masks)):\n",
        "    if i == plots:\n",
        "        break\n",
        "    img = img.to(device).unsqueeze(0)\n",
        "    predictions.append((model(img).detach().cpu()[0] > 0.5).float())\n",
        "    image_mask.append(mask)\n",
        "plotn(plots, (predictions, image_mask), only_mask=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "There are also some formal metrics to evaluate the performance, which you can read about [here](https://towardsdatascience.com/metrics-to-evaluate-your-semantic-segmentation-model-6bcb99639aa2). The easiest one to understand is **pixel accuracy** - a percentage of pixels classified correctly."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RU5KGWXaTbso"
      },
      "source": [
        "## U-Net\n",
        "\n",
        "SegNet architecture is very natural, but it is not the most accurate. Indeed, we first apply pyramid CNN architecture to the original image, which reduces the spatial accuracy of image features. Then, when we reconstruct the image, we cannot correctly reconstruct the pixel positions.\n",
        "\n",
        "This leads us to the idea of **skip connections** between convolution layers in encoder and decoder. This architecture is very common for semantic segmentation, and is called **U-Net**. Skip connections at each convolution level helps network not to lose information about features from original input at this level.\n",
        "\n",
        "We will use quite simple CNN architecture here, but U-Net can also use more complex encoder for feature extraction, such as ResNet-50.\n",
        "\n",
        "<img src=\"images/unet.png\" width=\"70%\">\n",
        "\n",
        "> Image from paper: Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. [U-Net: Convolutional networks for biomedical image segmentation.](https://arxiv.org/pdf/1505.04597.pdf)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:17:45.665392Z",
          "iopub.status.busy": "2022-04-08T16:17:45.665102Z",
          "iopub.status.idle": "2022-04-08T16:17:45.691051Z",
          "shell.execute_reply": "2022-04-08T16:17:45.690314Z",
          "shell.execute_reply.started": "2022-04-08T16:17:45.665341Z"
        },
        "id": "ZLKGrI4YTbs9",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class UNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.enc_conv0 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=(3,3), padding=1)\n",
        "        self.act0 = nn.ReLU()\n",
        "        self.bn0 = nn.BatchNorm2d(16)\n",
        "        self.pool0 = nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.enc_conv1 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3,3), padding=1)\n",
        "        self.act1 = nn.ReLU()\n",
        "        self.bn1 = nn.BatchNorm2d(32)\n",
        "        self.pool1 = nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.enc_conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3,3), padding=1)\n",
        "        self.act2 = nn.ReLU()\n",
        "        self.bn2 = nn.BatchNorm2d(64)\n",
        "        self.pool2 =  nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.enc_conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3,3), padding=1)\n",
        "        self.act3 = nn.ReLU()\n",
        "        self.bn3 = nn.BatchNorm2d(128)\n",
        "        self.pool3 =  nn.MaxPool2d(kernel_size=(2,2))\n",
        "\n",
        "        self.bottleneck_conv = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3,3), padding=1)\n",
        "        \n",
        "        self.upsample0 =  nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        self.dec_conv0 = nn.Conv2d(in_channels=384, out_channels=128, kernel_size=(3,3), padding=1)\n",
        "        self.dec_act0 = nn.ReLU()\n",
        "        self.dec_bn0 = nn.BatchNorm2d(128)\n",
        "\n",
        "        self.upsample1 =  nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        self.dec_conv1 =  nn.Conv2d(in_channels=192, out_channels=64, kernel_size=(3,3), padding=1)\n",
        "        self.dec_act1 = nn.ReLU()\n",
        "        self.dec_bn1 = nn.BatchNorm2d(64)\n",
        "\n",
        "        self.upsample2 = nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        self.dec_conv2 = nn.Conv2d(in_channels=96, out_channels=32, kernel_size=(3,3), padding=1)\n",
        "        self.dec_act2 = nn.ReLU()\n",
        "        self.dec_bn2 = nn.BatchNorm2d(32)\n",
        "\n",
        "        self.upsample3 = nn.UpsamplingBilinear2d(scale_factor=2)\n",
        "        self.dec_conv3 = nn.Conv2d(in_channels=48, out_channels=1, kernel_size=(1,1))\n",
        "\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "\n",
        "    def forward(self, x):\n",
        "        e0 = self.pool0(self.bn0(self.act0(self.enc_conv0(x))))\n",
        "        e1 = self.pool1(self.bn1(self.act1(self.enc_conv1(e0))))\n",
        "        e2 = self.pool2(self.bn2(self.act2(self.enc_conv2(e1))))\n",
        "        e3 = self.pool3(self.bn3(self.act3(self.enc_conv3(e2))))\n",
        "\n",
        "        cat0 = self.bn0(self.act0(self.enc_conv0(x)))\n",
        "        cat1 = self.bn1(self.act1(self.enc_conv1(e0)))\n",
        "        cat2 = self.bn2(self.act2(self.enc_conv2(e1)))\n",
        "        cat3 = self.bn3(self.act3(self.enc_conv3(e2)))\n",
        "\n",
        "        b = self.bottleneck_conv(e3)\n",
        "\n",
        "        d0 = self.dec_bn0(self.dec_act0(self.dec_conv0(torch.cat((self.upsample0(b), cat3), dim=1))))\n",
        "        d1 = self.dec_bn1(self.dec_act1(self.dec_conv1(torch.cat((self.upsample1(d0), cat2), dim=1))))\n",
        "        d2 = self.dec_bn2(self.dec_act2(self.dec_conv2(torch.cat((self.upsample2(d1), cat1), dim=1))))\n",
        "        d3 = self.sigmoid(self.dec_conv3(torch.cat((self.upsample3(d2), cat0), dim=1)))\n",
        "        return d3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T16:17:45.692880Z",
          "iopub.status.busy": "2022-04-08T16:17:45.692240Z",
          "iopub.status.idle": "2022-04-08T16:17:45.719635Z",
          "shell.execute_reply": "2022-04-08T16:17:45.719023Z",
          "shell.execute_reply.started": "2022-04-08T16:17:45.692842Z"
        },
        "id": "15GA_43BTbtI",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "model = UNet().to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
        "loss_fn = nn.BCEWithLogitsLoss()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "execution": {
          "iopub.execute_input": "2022-04-08T16:17:45.721443Z",
          "iopub.status.busy": "2022-04-08T16:17:45.721062Z",
          "iopub.status.idle": "2022-04-08T16:20:23.193420Z",
          "shell.execute_reply": "2022-04-08T16:20:23.191453Z",
          "shell.execute_reply.started": "2022-04-08T16:17:45.721410Z"
        },
        "id": "_dgiuvVVFMpr",
        "outputId": "438c570f-9480-48c6-bce6-14fbdf5b2d5f",
        "trusted": true
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 30/30 [29:07<00:00, 58.26s/it, train loss:=0.595, test loss:=0.572] \n"
          ]
        }
      ],
      "source": [
        "train(dataloaders, model, loss_fn, optimizer, epochs, device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 203
        },
        "id": "iEu5wjuMFMps",
        "outputId": "cf76869c-cf86-493f-fa73-a8790d21bf20"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 8392 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 7048 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "model.eval()\n",
        "predictions = []\n",
        "image_mask = []\n",
        "plots = 5\n",
        "images, masks = test_dataset[0], test_dataset[1]\n",
        "for i, (img, mask) in enumerate(zip(images, masks)):\n",
        "    if i == plots:\n",
        "        break\n",
        "    img = img.to(device).unsqueeze(0)\n",
        "    predictions.append((model(img).detach().cpu()[0] > 0.5).float())\n",
        "    image_mask.append(mask)\n",
        "plotn(plots, (predictions, image_mask), only_mask=True)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "SemanticSegmentation.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
