UNET.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/AriPathak/ResUnet/blob/main/UNET.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KU69kGdVwoDn"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.transforms.functional as TF"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m9uuFPmE-6Kn"
      },
      "outputs": [],
      "source": [
        "class Rescale(object):\n",
        "    def __init__(self, output_size):\n",
        "        assert isinstance(output_size, (int, tuple))\n",
        "        self.output_size = output_size\n",
        "\n",
        "    def __call__(self, sample):\n",
        "        image = sample\n",
        "        h, w = image.size\n",
        "        if h > w:\n",
        "            new_h, new_w = self.output_size * h / w, self.output_size\n",
        "        else:\n",
        "            new_h, new_w = self.output_size, self.output_size * w / h\n",
        "\n",
        "        new_h, new_w = int(new_h), int(new_w)\n",
        "\n",
        "        img = TF.resize(image, (new_w, new_h))\n",
        "\n",
        "        return img"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MNoDLWer_Kj-",
        "outputId": "74bdace3-32b3-4d3a-a734-e22470d4f8f2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading https://thor.robots.ox.ac.uk/datasets/pets/images.tar.gz to oxford-iiit-pet/images.tar.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 791918971/791918971 [00:38<00:00, 20610736.67it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting oxford-iiit-pet/images.tar.gz to oxford-iiit-pet\n",
            "Downloading https://thor.robots.ox.ac.uk/datasets/pets/annotations.tar.gz to oxford-iiit-pet/annotations.tar.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 19173078/19173078 [00:01<00:00, 9919986.26it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting oxford-iiit-pet/annotations.tar.gz to oxford-iiit-pet\n"
          ]
        }
      ],
      "source": [
        "num_classes = 3 # as specified above, there are 3 classes for each pixel\n",
        "img_dim = 128\n",
        "trainset = torchvision.datasets.OxfordIIITPet(root=\"./\", download=True, target_types='segmentation',\n",
        "                                              transform=transforms.Compose([\n",
        "                                                            Rescale(img_dim),\n",
        "                                                            transforms.CenterCrop(img_dim),\n",
        "                                                            transforms.ToTensor()]),\n",
        "                                              target_transform=transforms.Compose([\n",
        "                                                            Rescale(img_dim),\n",
        "                                                            transforms.CenterCrop(img_dim),\n",
        "                                                            transforms.Lambda(lambda x: torch.from_numpy(np.array(x) - 1).long())]))\n",
        "testset = torchvision.datasets.OxfordIIITPet(root=\"./\", download=True, target_types='segmentation', split='test',\n",
        "                                             transform=transforms.Compose([\n",
        "                                                            Rescale(img_dim),\n",
        "                                                            transforms.CenterCrop(img_dim),\n",
        "                                                            transforms.ToTensor()]),\n",
        "                                             target_transform=transforms.Compose([\n",
        "                                                            Rescale(img_dim),\n",
        "                                                            transforms.CenterCrop(img_dim),\n",
        "                                                            transforms.Lambda(lambda x: torch.from_numpy(np.array(x) - 1).long())]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9gkr2DxUxBCs"
      },
      "outputs": [],
      "source": [
        "from torchvision.models import resnext50_32x4d\n",
        "\n",
        "class ConvRelu(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, kernel, padding):\n",
        "        super().__init__()\n",
        "\n",
        "        self.convrelu = nn.Sequential(\n",
        "            nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n",
        "            nn.ReLU(inplace=True)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.convrelu(x)\n",
        "        return x\n",
        "\n",
        "class DecoderBlock(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels):\n",
        "        super().__init__()\n",
        "\n",
        "        self.conv1 = ConvRelu(in_channels, in_channels // 4, 1, 0)\n",
        "\n",
        "        self.deconv = nn.ConvTranspose2d(in_channels // 4, in_channels // 4, kernel_size=4,\n",
        "                                          stride=2, padding=1, output_padding=0)\n",
        "\n",
        "        self.conv2 = ConvRelu(in_channels // 4, out_channels, 1, 0)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv1(x)\n",
        "        x = self.deconv(x)\n",
        "        x = self.conv2(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "class ResUnetDecoderBlock(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels):\n",
        "        super().__init__()\n",
        "\n",
        "        self.conv1 = ConvRelu(in_channels, in_channels, 3, 1)\n",
        "\n",
        "        self.deconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=3,\n",
        "                                          stride=2, padding=1, output_padding=1)\n",
        "\n",
        "        self.conv2 = ConvRelu(out_channels, out_channels, 1, 0)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv1(x)\n",
        "        x = self.deconv(x)\n",
        "        x = self.conv2(x)\n",
        "\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fQ861jTVxjEA"
      },
      "outputs": [],
      "source": [
        "class ResBlock(nn.Module):\n",
        "  def __init__(self, in_channels, out_channels, stride=1):\n",
        "    super().__init__()\n",
        "    self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride = 1, padding=1)\n",
        "    self.conv2 = nn.Conv2d(out_channels, out_channels, 3, stride, padding=1)\n",
        "    self.norm1 = nn.BatchNorm2d(out_channels)\n",
        "    self.norm2 = nn.BatchNorm2d(out_channels)\n",
        "    self.shortcut = nn.Sequential()\n",
        "    if stride != 1 or in_channels != out_channels:\n",
        "      self.shortcut = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, stride),\n",
        "                                        nn.BatchNorm2d(out_channels))\n",
        "  def forward(self, x):\n",
        "    res = x\n",
        "    x = F.relu(self.norm1(self.conv1(x)))\n",
        "    x = self.norm2(self.conv2(x))\n",
        "    res = F.relu(self.shortcut(res))\n",
        "    x += res\n",
        "    return F.relu(x)\n",
        "\n",
        "\n",
        "class ResUNet(nn.Module):\n",
        "  def __init__(self, num_classes, num_channels):\n",
        "    super().__init__()\n",
        "    self.e0 = nn.Conv2d(num_channels, 32, 3, padding=1)\n",
        "    self.bn1 = nn.BatchNorm2d(32)\n",
        "    self.e1 = ResBlock(32, 64, stride=1) #why is stride 1 here ari?\n",
        "    self.e2 = ResBlock(64, 128, stride=2)\n",
        "    self.e3 = ResBlock(128, 256, stride=2)\n",
        "    self.e4 = ResBlock(256, 512, stride=2)\n",
        "\n",
        "\n",
        "    self.decoder4 = ResUnetDecoderBlock(512, 256)\n",
        "    self.decoder3 = ResUnetDecoderBlock(256, 128)\n",
        "    self.decoder2 = ResUnetDecoderBlock(128, 64)\n",
        "\n",
        "    self.last_conv0 = nn.Conv2d(64, num_classes, (1,1), padding=0)\n",
        "\n",
        "  def forward(self, x):\n",
        "    res0 = F.relu(self.bn1(self.e0(x)))\n",
        "    res1 = self.e1(res0)\n",
        "    res2 = self.e2(res1)\n",
        "    res3 = self.e3(res2)\n",
        "    res4 = self.e4(res3)\n",
        "\n",
        "    x = self.decoder4(res4) + res3\n",
        "    x = self.decoder3(x) + res2\n",
        "    x = self.decoder2(x) + res1\n",
        "    out = self.last_conv0(x)\n",
        "    return F.softmax(out, dim=1)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gYDuWC793O2N"
      },
      "outputs": [],
      "source": [
        "class Unet(nn.Module):\n",
        "  def __init__(self, n_classes):\n",
        "    super().__init__()\n",
        "\n",
        "    self.base_model = resnext50_32x4d(pretrained=True)\n",
        "    self.base_layers = list(self.base_model.children())\n",
        "    filters = [4*64, 4*128, 4*256, 4*512]\n",
        "\n",
        "    self.encoder0 = nn.Sequential(*self.base_layers[:3])\n",
        "    self.encoder1 = nn.Sequential(*self.base_layers[4])\n",
        "    self.encoder2 = nn.Sequential(*self.base_layers[5])\n",
        "    self.encoder3 = nn.Sequential(*self.base_layers[6])\n",
        "    self.encoder4 = nn.Sequential(*self.base_layers[7])\n",
        "\n",
        "    self.decoder4 = DecoderBlock(filters[3], filters[2])\n",
        "    self.decoder3 = DecoderBlock(filters[2], filters[1])\n",
        "    self.decoder2 = DecoderBlock(filters[1], filters[0])\n",
        "    self.decoder1 = DecoderBlock(filters[0], filters[0])\n",
        "\n",
        "    self.last_conv1 = nn.Conv2d(256, n_classes, 1, padding=0)\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.encoder0(x)\n",
        "    e1 = self.encoder1(x)\n",
        "    e2 = self.encoder2(e1)\n",
        "    e3 = self.encoder3(e2)\n",
        "    e4 = self.encoder4(e3)\n",
        "\n",
        "    d4 = self.decoder4(e4) + e3\n",
        "    d3 = self.decoder3(d4) + e2\n",
        "    d2 = self.decoder2(d3) + e1\n",
        "    d1 = self.decoder1(d2)\n",
        "\n",
        "    out = self.last_conv1(d1)\n",
        "    return F.softmax(out, dim=1)\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p3DgPH6HYLQV"
      },
      "outputs": [],
      "source": [
        "class mIoULoss(nn.Module):\n",
        "    def __init__(self, weight=None, size_average=True, n_classes=2):\n",
        "        super(mIoULoss, self).__init__()\n",
        "        self.classes = n_classes\n",
        "\n",
        "    def to_one_hot(self, tensor):\n",
        "        n,h,w = tensor.size()\n",
        "        one_hot = torch.zeros(n,self.classes,h,w).to(tensor.device).scatter_(1,tensor.view(n,1,h,w),1)\n",
        "        return one_hot\n",
        "\n",
        "    def forward(self, inputs, target):\n",
        "        # inputs => N x Classes x H x W\n",
        "        # target_oneHot => N x Classes x H x W\n",
        "\n",
        "        N = inputs.size()[0]\n",
        "\n",
        "        # predicted probabilities for each pixel along channel\n",
        "        #inputs = F.softmax(inputs,dim=1)\n",
        "\n",
        "        # Numerator Product\n",
        "        target_oneHot = self.to_one_hot(target)\n",
        "        inter = inputs * target_oneHot\n",
        "        ## Sum over all pixels N x C x H x W => N x C\n",
        "        inter = inter.view(N,self.classes,-1).sum(2)\n",
        "\n",
        "        #Denominator\n",
        "        union= inputs + target_oneHot - (inputs*target_oneHot)\n",
        "        ## Sum over all pixels N x C x H x W => N x C\n",
        "        union = union.view(N,self.classes,-1).sum(2)\n",
        "\n",
        "        loss = inter/union\n",
        "\n",
        "        ## Return average loss over classes and batch\n",
        "        return 1-loss.mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aCEva9pA3O8t"
      },
      "outputs": [],
      "source": [
        "device = 'cuda:0'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sgScewe4ZUvK"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "from tqdm.notebook import tqdm\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "def training_loop(model, num_epochs, batch_size, trainset, criterion, optimizer, scheduler):\n",
        "    train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True)\n",
        "    train_loss = []\n",
        "    model.cuda()\n",
        "    for epoch in range(1, num_epochs + 1):\n",
        "        batches = tqdm(train_loader, initial=1)\n",
        "        model.train()\n",
        "        running_loss = 0\n",
        "        for images, labels in batches:\n",
        "            images = images.to(device)\n",
        "            labels = labels.to(device).squeeze()\n",
        "            predictions = model(images)\n",
        "            optimizer.zero_grad()\n",
        "            loss = criterion(predictions, labels)\n",
        "            #print(loss.item())\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            running_loss += loss.item()\n",
        "        if scheduler:\n",
        "            scheduler.step()\n",
        "        #epoch_iou = mean_iou(trainset, model)\n",
        "        epoch_iou = 0\n",
        "        train_loss.append(running_loss / len(train_loader))\n",
        "        print(\"Epoch: {0:02d} | Training Loss: {1:.5f} | IOU: {2:.5f}\".format(epoch, train_loss[-1], epoch_iou))\n",
        "    return train_loss\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TUdrjcIIaVTZ"
      },
      "outputs": [],
      "source": [
        "import torch.optim as optim\n",
        "\n",
        "learning_rate = 1e-3\n",
        "n_epochs = 20\n",
        "batch_size = 16\n",
        "u_net = ResUNet(3, 3)\n",
        "\n",
        "criterion = mIoULoss(n_classes=3)\n",
        "optimizer = optim.Adam(u_net.parameters(), lr=learning_rate)\n",
        "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)\n",
        "training_loop(u_net, n_epochs, batch_size, trainset, criterion, optimizer, scheduler)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gu0kQQ3WMCPn"
      },
      "outputs": [],
      "source": [
        "torch.save(u_net.state_dict(), \"UNET_RESBACKBONE_OxfordIII.pth\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PA-vOR46BVhM"
      },
      "outputs": [],
      "source": [
        "import matplotlib.gridspec as gs\n",
        "\n",
        "\n",
        "num_samples = 3\n",
        "img_grid = gs.GridSpec(num_samples, 3, width_ratios=[1, 1, 1], wspace=0.2, hspace=0.2)\n",
        "fig = plt.figure(figsize=(10, 30))\n",
        "\n",
        "for idx, i in enumerate(np.random.choice(len(trainset), num_samples)):\n",
        "    image, label = trainset[i]\n",
        "    image = image.to(device)\n",
        "    pred = u_net(image.unsqueeze(dim=0)).squeeze()\n",
        "    print(pred.shape)\n",
        "    pred = torch.argmax(pred, dim=0).cpu() #dim 0 is the output class channels/3 rgb maps\n",
        "    image = image.cpu()\n",
        "    plt.subplot(img_grid[idx, 0]).imshow(image.permute(1, 2, 0))\n",
        "    plt.subplot(img_grid[idx, 1]).imshow(label.squeeze())\n",
        "    plt.subplot(img_grid[idx, 2]).imshow(pred.squeeze())"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}