B3_PSPNet_v2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BQNLZ86t9R3G"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Kiểm tra GPU được cấp phát</h3>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GOVtzLUjjZ1W"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vigGQG1--GdX"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Cài đặt bổ sung một số thư viện</h3>\n",
        "</div>\n",
        "Nền tảng Google Colab cung cấp môi trường với các thư viện Machine Learning, Deep Learning cơ bản đã được cài đặt sẵn, phần này sẽ cài đặt một số thư viện sử dụng thêm."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tChf3YMRYr4m"
      },
      "outputs": [],
      "source": [
        "!pip install segmentation_models_pytorch\n",
        "!pip install albumentations\n",
        "!pip install --upgrade gdown"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7dnqP5oJ-yEs"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3></h3>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n2ZRXglT_5Rm"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Import thư viện</h3>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qCnMyzJgZqJr"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "import segmentation_models_pytorch as smp\n",
        "import albumentations as A\n",
        "from albumentations.pytorch import ToTensorV2 # np.array -> torch.tensor\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "from glob import glob\n",
        "from torchvision.datasets import VOCSegmentation"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Xz7v36fKCFqV"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Định nghĩa Dataset</h3>\n",
        "</div>\n",
        "Viết class kế thừa từ class Dataset cung cấp sẵn trong PyTorch để đọc dữ liệu từ ổ cứng. Yêu cầu viết đủ 3 hàm __init__() để khởi tạo class, __len__() để trả về số điểm dữ liệu có trong tập dữ liệu và __getitem__() trả về 1 điểm dữ liệu cụ thể. Trong phần này, do tập dữ liệu PASCAL VOC đã rất phổ biến nên sẽ tận dụng Class Dataset đã được viết sẵn. Tham khảo thêm: https://albumentations.ai/docs/autoalbument/examples/pascal_voc/\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GNv3AOficqQK"
      },
      "outputs": [],
      "source": [
        "cv2.setNumThreads(0)\n",
        "cv2.ocl.setUseOpenCL(False)\n",
        "\n",
        "VOC_CLASSES = [\n",
        "    \"background\",\n",
        "    \"aeroplane\",\n",
        "    \"bicycle\",\n",
        "    \"bird\",\n",
        "    \"boat\",\n",
        "    \"bottle\",\n",
        "    \"bus\",\n",
        "    \"car\",\n",
        "    \"cat\",\n",
        "    \"chair\",\n",
        "    \"cow\",\n",
        "    \"diningtable\",\n",
        "    \"dog\",\n",
        "    \"horse\",\n",
        "    \"motorbike\",\n",
        "    \"person\",\n",
        "    \"potted plant\",\n",
        "    \"sheep\",\n",
        "    \"sofa\",\n",
        "    \"train\",\n",
        "    \"tv/monitor\",\n",
        "]\n",
        "\n",
        "VOC_COLORMAP = [\n",
        "    [0, 0, 0],\n",
        "    [128, 0, 0],\n",
        "    [0, 128, 0],\n",
        "    [128, 128, 0],\n",
        "    [0, 0, 128],\n",
        "    [128, 0, 128],\n",
        "    [0, 128, 128],\n",
        "    [128, 128, 128],\n",
        "    [64, 0, 0],\n",
        "    [192, 0, 0],\n",
        "    [64, 128, 0],\n",
        "    [192, 128, 0],\n",
        "    [64, 0, 128],\n",
        "    [192, 0, 128],\n",
        "    [64, 128, 128],\n",
        "    [192, 128, 128],\n",
        "    [0, 64, 0],\n",
        "    [128, 64, 0],\n",
        "    [0, 192, 0],\n",
        "    [128, 192, 0],\n",
        "    [0, 64, 128],\n",
        "]\n",
        "\n",
        "class PascalVOCSearchDataset(VOCSegmentation):\n",
        "    def __init__(self, root=\"~/data/pascal_voc\", image_set=\"train\", download=True, transform=None):\n",
        "        super().__init__(root=root, image_set=image_set, download=download, transform=transform)\n",
        "\n",
        "    @staticmethod\n",
        "    def _convert_to_segmentation_mask(mask):\n",
        "        # This function converts a mask from the Pascal VOC format to the format required by AutoAlbument.\n",
        "        #\n",
        "        # Pascal VOC uses an RGB image to encode the segmentation mask for that image. RGB values of a pixel\n",
        "        # encode the pixel's class.\n",
        "        #\n",
        "        # AutoAlbument requires a segmentation mask to be a NumPy array with the shape [height, width, num_classes].\n",
        "        # Each channel in this mask should encode values for a single class. Pixel in a mask channel should have\n",
        "        # a value of 1.0 if the pixel of the image belongs to this class and 0.0 otherwise.\n",
        "        height, width = mask.shape[:2]\n",
        "        segmentation_mask = np.zeros((height, width, len(VOC_COLORMAP)), dtype=np.float32)\n",
        "        for label_index, label in enumerate(VOC_COLORMAP):\n",
        "            segmentation_mask[:, :, label_index] = np.all(mask == label, axis=-1).astype(float)\n",
        "        return segmentation_mask #0, 1, 2, 3, ..., 20 (H, W, C) -> (H, W, 1) -> (H, W) #numpy\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        image = cv2.imread(self.images[index])\n",
        "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "        mask = cv2.imread(self.masks[index])\n",
        "        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)\n",
        "        mask = self._convert_to_segmentation_mask(mask)\n",
        "        if self.transform is not None:\n",
        "            transformed = self.transform(image=image, mask=mask)\n",
        "            image = transformed[\"image\"]\n",
        "            mask = transformed[\"mask\"]\n",
        "        return image, mask.argmax(dim=2).squeeze() #torch.tensor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "en9W5v_xChlD"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Định nghĩa các phép augmentation trên ảnh</h3>\n",
        "</div>\n",
        "Sử dụng thư viện Albumentations, tham khảo thêm: https://albumentations.ai/docs/api_reference/full_reference/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e4RPW6nMgx4a"
      },
      "outputs": [],
      "source": [
        "trainsize = 256\n",
        "\n",
        "train_transform = A.Compose([\n",
        "    A.Resize(width=trainsize, height=trainsize),\n",
        "    A.HorizontalFlip(),\n",
        "    A.RandomBrightnessContrast(),\n",
        "    A.Blur(),\n",
        "    A.Sharpen(),\n",
        "    A.RGBShift(),\n",
        "    A.Cutout(num_holes=5, max_h_size=25, max_w_size=25, fill_value=0),\n",
        "    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0),\n",
        "    ToTensorV2(),\n",
        "])\n",
        "\n",
        "test_trainsform = A.Compose([\n",
        "    A.Resize(width=trainsize, height=trainsize),\n",
        "    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0),\n",
        "    ToTensorV2(), # numpy.array -> torch.tensor (B, 3, H, W)\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r81wtP0PCzUD"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Đoạn code dùng để convert ảnh sau khi đã chuẩn hoá thành ảnh ban đầu</h3>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RSdnanKkjlD3"
      },
      "outputs": [],
      "source": [
        "class UnNormalize(object):\n",
        "    def __init__(self, mean, std):\n",
        "        self.mean = mean\n",
        "        self.std = std\n",
        "        \n",
        "    def __call__(self, tensor):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.\n",
        "        Returns:\n",
        "            Tensor: Normalized image.\n",
        "        \"\"\"\n",
        "        for t, m, s in zip(tensor, self.mean, self.std):\n",
        "            t.mul_(s).add_(m)\n",
        "            # The normalize code -> t.sub_(m).div_(s)\n",
        "        return tensor\n",
        "    \n",
        "unorm = UnNormalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x_CbSVe6C_bo"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Kiểm tra 1 cặp ảnh đầu vào và ảnh kết quả phân vùng trước khi đưa vào mô hình training</h3>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "06Kt7pRTfD8H"
      },
      "outputs": [],
      "source": [
        "train_dataset = PascalVOCSearchDataset(image_set=\"train\", download=True, transform=train_transform)\n",
        "test_dataset = PascalVOCSearchDataset(image_set=\"val\", download=False, transform=test_trainsform)\n",
        "\n",
        "image, mask = train_dataset.__getitem__(10)\n",
        "plt.subplot(1, 2, 1)\n",
        "plt.imshow(unorm(image).permute(1, 2, 0))\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.imshow(mask)\n",
        "plt.show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "XNubRiEKDHg0"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Backbone ResNet cho mô hình PSPNet</h3>\n",
        "</div>\n",
        "Tham khảo: https://github.com/hszhao/semseg/tree/master/model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3nSSXeMMmXy5"
      },
      "outputs": [],
      "source": [
        "import math\n",
        "import torch.utils.model_zoo as model_zoo\n",
        "\n",
        "\n",
        "__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',\n",
        "           'resnet152']\n",
        "\n",
        "\n",
        "model_urls = {\n",
        "    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n",
        "    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n",
        "    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n",
        "    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n",
        "    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n",
        "}\n",
        "\n",
        "\n",
        "def conv3x3(in_planes, out_planes, stride=1):\n",
        "    \"\"\"3x3 convolution with padding\"\"\"\n",
        "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
        "                     padding=1, bias=False)\n",
        "\n",
        "\n",
        "class BasicBlock(nn.Module):\n",
        "    expansion = 1\n",
        "\n",
        "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
        "        super(BasicBlock, self).__init__()\n",
        "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
        "        self.bn1 = nn.BatchNorm2d(planes)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.conv2 = conv3x3(planes, planes)\n",
        "        self.bn2 = nn.BatchNorm2d(planes)\n",
        "        self.downsample = downsample\n",
        "        self.stride = stride\n",
        "\n",
        "    def forward(self, x):\n",
        "        residual = x\n",
        "\n",
        "        out = self.conv1(x)\n",
        "        out = self.bn1(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv2(out)\n",
        "        out = self.bn2(out)\n",
        "\n",
        "        if self.downsample is not None:\n",
        "            residual = self.downsample(x)\n",
        "\n",
        "        out += residual\n",
        "        out = self.relu(out)\n",
        "\n",
        "        return out\n",
        "\n",
        "\n",
        "class Bottleneck(nn.Module):\n",
        "    expansion = 4\n",
        "\n",
        "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
        "        super(Bottleneck, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(planes)\n",
        "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n",
        "                               padding=1, bias=False)\n",
        "        self.bn2 = nn.BatchNorm2d(planes)\n",
        "        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n",
        "        self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.downsample = downsample\n",
        "        self.stride = stride\n",
        "\n",
        "    def forward(self, x):\n",
        "        residual = x\n",
        "\n",
        "        out = self.conv1(x)\n",
        "        out = self.bn1(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv2(out)\n",
        "        out = self.bn2(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv3(out)\n",
        "        out = self.bn3(out)\n",
        "\n",
        "        if self.downsample is not None:\n",
        "            residual = self.downsample(x)\n",
        "\n",
        "        out += residual\n",
        "        out = self.relu(out)\n",
        "\n",
        "        return out\n",
        "\n",
        "\n",
        "class ResNet(nn.Module):\n",
        "\n",
        "    def __init__(self, block, layers, num_classes=1000, deep_base=True):\n",
        "        super(ResNet, self).__init__()\n",
        "        self.deep_base = deep_base\n",
        "        if not self.deep_base:\n",
        "            self.inplanes = 64\n",
        "            self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
        "            self.bn1 = nn.BatchNorm2d(64)\n",
        "        else:\n",
        "            self.inplanes = 128\n",
        "            self.conv1 = conv3x3(3, 64, stride=2)\n",
        "            self.bn1 = nn.BatchNorm2d(64)\n",
        "            self.conv2 = conv3x3(64, 64)\n",
        "            self.bn2 = nn.BatchNorm2d(64)\n",
        "            self.conv3 = conv3x3(64, 128)\n",
        "            self.bn3 = nn.BatchNorm2d(128)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
        "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
        "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
        "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n",
        "        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n",
        "        self.avgpool = nn.AvgPool2d(7, stride=1)\n",
        "        self.fc = nn.Linear(512 * block.expansion, num_classes)\n",
        "\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.Conv2d):\n",
        "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
        "            elif isinstance(m, nn.BatchNorm2d):\n",
        "                nn.init.constant_(m.weight, 1)\n",
        "                nn.init.constant_(m.bias, 0)\n",
        "\n",
        "    def _make_layer(self, block, planes, blocks, stride=1):\n",
        "        downsample = None\n",
        "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
        "            downsample = nn.Sequential(\n",
        "                nn.Conv2d(self.inplanes, planes * block.expansion,\n",
        "                          kernel_size=1, stride=stride, bias=False),\n",
        "                nn.BatchNorm2d(planes * block.expansion),\n",
        "            )\n",
        "\n",
        "        layers = []\n",
        "        layers.append(block(self.inplanes, planes, stride, downsample))\n",
        "        self.inplanes = planes * block.expansion\n",
        "        for i in range(1, blocks):\n",
        "            layers.append(block(self.inplanes, planes))\n",
        "\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.relu(self.bn1(self.conv1(x)))\n",
        "        if self.deep_base:\n",
        "            x = self.relu(self.bn2(self.conv2(x)))\n",
        "            x = self.relu(self.bn3(self.conv3(x)))\n",
        "        x = self.maxpool(x)\n",
        "\n",
        "        x = self.layer1(x)\n",
        "        x = self.layer2(x)\n",
        "        x = self.layer3(x)\n",
        "        x = self.layer4(x)\n",
        "\n",
        "        x = self.avgpool(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = self.fc(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "def resnet18(pretrained=True, **kwargs):\n",
        "    \"\"\"Constructs a ResNet-18 model.\n",
        "    Args:\n",
        "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
        "    \"\"\"\n",
        "    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)\n",
        "    if pretrained:\n",
        "        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))\n",
        "    return model\n",
        "\n",
        "\n",
        "def resnet34(pretrained=True, **kwargs):\n",
        "    \"\"\"Constructs a ResNet-34 model.\n",
        "    Args:\n",
        "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
        "    \"\"\"\n",
        "    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)\n",
        "    if pretrained:\n",
        "        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))\n",
        "    return model\n",
        "\n",
        "\n",
        "def resnet50(pretrained=True, **kwargs):\n",
        "    \"\"\"Constructs a ResNet-50 model.\n",
        "    Args:\n",
        "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
        "    \"\"\"\n",
        "    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)\n",
        "    if pretrained:\n",
        "        # model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))\n",
        "        model_path = 'resnet50_v2.pth'\n",
        "        model.load_state_dict(torch.load(model_path), strict=False)\n",
        "    return model\n",
        "\n",
        "\n",
        "def resnet101(pretrained=True, **kwargs):\n",
        "    \"\"\"Constructs a ResNet-101 model.\n",
        "    Args:\n",
        "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
        "    \"\"\"\n",
        "    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)\n",
        "    if pretrained:\n",
        "        # model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))\n",
        "        model_path = 'resnet101_v2.pth'\n",
        "        model.load_state_dict(torch.load(model_path), strict=False)\n",
        "    return model\n",
        "\n",
        "\n",
        "def resnet152(pretrained=True, **kwargs):\n",
        "    \"\"\"Constructs a ResNet-152 model.\n",
        "    Args:\n",
        "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
        "    \"\"\"\n",
        "    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)\n",
        "    if pretrained:\n",
        "        # model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))\n",
        "        model_path = './initmodel/resnet152_v2.pth'\n",
        "        model.load_state_dict(torch.load(model_path), strict=False)\n",
        "    return model"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Mô hình PSPNet</h3>\n",
        "</div>\n",
        "Tham khảo: https://github.com/hszhao/semseg/tree/master/model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "class PPM(nn.Module):\n",
        "    def __init__(self, in_dim, reduction_dim, bins):\n",
        "        super(PPM, self).__init__()\n",
        "        self.features = []\n",
        "        for bin in bins:\n",
        "            self.features.append(nn.Sequential(\n",
        "                nn.AdaptiveAvgPool2d(bin),\n",
        "                nn.Conv2d(in_dim, reduction_dim, kernel_size=1, bias=False),\n",
        "                nn.BatchNorm2d(reduction_dim),\n",
        "                nn.ReLU(inplace=True)\n",
        "            ))\n",
        "        self.features = nn.ModuleList(self.features)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x_size = x.size()\n",
        "        out = [x]\n",
        "        for f in self.features:\n",
        "            out.append(F.interpolate(f(x), x_size[2:], mode='bilinear', align_corners=True))\n",
        "        return torch.cat(out, 1) \n",
        "\n",
        "class PSPNet(nn.Module):\n",
        "    def __init__(self, layers=50, bins=(1, 2, 3, 6), dropout=0.1, classes=2, zoom_factor=8, use_ppm=True, criterion=nn.CrossEntropyLoss(ignore_index=255), pretrained=True):\n",
        "        super(PSPNet, self).__init__()\n",
        "        assert layers in [50, 101, 152]\n",
        "        assert 2048 % len(bins) == 0\n",
        "        assert classes > 1\n",
        "        assert zoom_factor in [1, 2, 4, 8]\n",
        "        self.zoom_factor = zoom_factor\n",
        "        self.use_ppm = use_ppm\n",
        "        self.criterion = criterion\n",
        "\n",
        "        if layers == 50:\n",
        "            resnet = resnet50(pretrained=pretrained)\n",
        "        elif layers == 101:\n",
        "            resnet = resnet101(pretrained=pretrained)\n",
        "        else:\n",
        "            resnet = resnet152(pretrained=pretrained)\n",
        "        self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.conv2, resnet.bn2, resnet.relu, resnet.conv3, resnet.bn3, resnet.relu, resnet.maxpool)\n",
        "        self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4\n",
        "\n",
        "        for n, m in self.layer3.named_modules():\n",
        "            if 'conv2' in n:\n",
        "                m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)\n",
        "            elif 'downsample.0' in n:\n",
        "                m.stride = (1, 1)\n",
        "        for n, m in self.layer4.named_modules():\n",
        "            if 'conv2' in n:\n",
        "                m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)\n",
        "            elif 'downsample.0' in n:\n",
        "                m.stride = (1, 1)\n",
        "\n",
        "        fea_dim = 2048\n",
        "        if use_ppm:\n",
        "            self.ppm = PPM(fea_dim, int(fea_dim/len(bins)), bins)\n",
        "            fea_dim *= 2\n",
        "        self.cls = nn.Sequential(\n",
        "            nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(512),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Dropout2d(p=dropout),\n",
        "            nn.Conv2d(512, classes, kernel_size=1)\n",
        "        )\n",
        "        if self.training:\n",
        "            self.aux = nn.Sequential(\n",
        "                nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),\n",
        "                nn.BatchNorm2d(256),\n",
        "                nn.ReLU(inplace=True),\n",
        "                nn.Dropout2d(p=dropout),\n",
        "                nn.Conv2d(256, classes, kernel_size=1)\n",
        "            )\n",
        "\n",
        "    def forward(self, x, y=None):\n",
        "        x_size = x.size()\n",
        "        assert (x_size[2]-1) % 8 == 0 and (x_size[3]-1) % 8 == 0\n",
        "        h = int((x_size[2] - 1) / 8 * self.zoom_factor + 1)\n",
        "        w = int((x_size[3] - 1) / 8 * self.zoom_factor + 1)\n",
        "\n",
        "        x = self.layer0(x)\n",
        "        x = self.layer1(x)\n",
        "        x = self.layer2(x)\n",
        "        x_tmp = self.layer3(x)\n",
        "        x = self.layer4(x_tmp)\n",
        "        if self.use_ppm:\n",
        "            x = self.ppm(x)\n",
        "        x = self.cls(x)\n",
        "        if self.zoom_factor != 1:\n",
        "            x = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)\n",
        "\n",
        "        if self.training:\n",
        "            aux = self.aux(x_tmp)\n",
        "            if self.zoom_factor != 1:\n",
        "                aux = F.interpolate(aux, size=(h, w), mode='bilinear', align_corners=True)\n",
        "            main_loss = self.criterion(x, y)\n",
        "            aux_loss = self.criterion(aux, y)\n",
        "            return x.max(1)[1], main_loss, aux_loss\n",
        "        else:\n",
        "            return x"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Download Pretrained mô hình cho backbone ResNet</h3>\n",
        "</div>\n",
        "Tham khảo: https://github.com/hszhao/semseg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!gdown 1w5pRmLJXvmQQA5PtCbHhZc_uC4o0YbmA\n",
        "!gdown 1V-sfLnqSuwTPgNMrqUp4HDZpqtoBH4xm\n",
        "!gdown 1pVzjWA1C-y4TniEkc06ygRO6e4Mec0wW"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-OZbtvpnD1ho"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Tạo AverageMeter</h3>\n",
        "</div>\n",
        "AverageMeter có nhiệm vụ lưu lại giá trị trung bình của độ chính xác, giá trị hàm loss, ... trong suốt quá trình training. Tham khảo thêm: https://discuss.pytorch.org/t/meaning-of-parameters/10655"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OJ0zgMSXy-7Z"
      },
      "outputs": [],
      "source": [
        "class AverageMeter(object):\n",
        "    def __init__(self):\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.val = 0\n",
        "        self.avg = 0\n",
        "        self.sum = 0\n",
        "        self.count = 0\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        self.val = val\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = self.sum / self.count"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "s8u9IBDlEQyF"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Lập trình hàm tính toán độ chính xác dựa trên IoU và Dice</h3>\n",
        "</div>\n",
        "Tham khảo: https://github.com/hszhao/semseg/blob/master/util/util.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9S0zUD210VJ5"
      },
      "outputs": [],
      "source": [
        "#metrics\n",
        "def intersectionAndUnionGPU(output, target, K, ignore_index=255):\n",
        "    # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1.\n",
        "    assert (output.dim() in [1, 2, 3])\n",
        "    assert output.shape == target.shape\n",
        "    output = output.view(-1)\n",
        "    target = target.view(-1)\n",
        "    output[target == ignore_index] = ignore_index\n",
        "    intersection = output[output == target]\n",
        "    area_intersection = torch.histc(intersection, bins=K, min=0, max=K-1)\n",
        "    area_output = torch.histc(output, bins=K, min=0, max=K-1)\n",
        "    area_target = torch.histc(target, bins=K, min=0, max=K-1)\n",
        "    area_union = area_output + area_target - area_intersection\n",
        "    return area_intersection, area_union, area_target "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MxkzTm38EXCo"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Chuẩn bị cho quá trình training</h3>\n",
        "</div>\n",
        "\n",
        "\n",
        "1.   Lựa chọn device: PyTorch yêu cầu lựa chọn cụ thể device sẽ train và yêu cầu người dùng tự move dữ liệu, mô hình vào device đã lựa chọn. Device có thể là \"cuda\" - tức là GPU NVIDIA hoặc \"cpu\".\n",
        "2.   Định nghĩa DataLoader, khác với Dataset là cách đọc dữ liệu từ ổ cứng, DataLoader ghép nhiều điểm dữ liệu vào cùng nhau tạo thành 1 batch để đưa vào train mô hình. Lưu ý thêm: batch_size nên đặt là 4, 8, 16, 32, ... và nên để lớn nhất có thể\n",
        "3.   Khởi tạo mô hình\n",
        "4.   Khởi tạo hàm loss\n",
        "5.   Khởi tạo thuật toán tối ưu (optimizer)\n",
        "6.   Khởi tạo các độ đo sẽ sử dụng để đánh giá hiệu năng của mô hình. Phần này sẽ sử dụng các hàm độ đo Dice và IoU được lập trình sẵn trong thư viện torchmetrics\n",
        "7.   Khởi tạo từng AverageMeter để lưu lại giá trị của từng độ đo, giá trị hàm loss, thời gian train, ... trong suốt quá trình train\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xldeI02fubce"
      },
      "outputs": [],
      "source": [
        "#device\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "#load data\n",
        "batch_size = 16\n",
        "n_workers = os.cpu_count()\n",
        "print(\"num_workers =\", n_workers)\n",
        "trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,\n",
        "                                          shuffle=True, num_workers=n_workers)\n",
        "testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,\n",
        "                                          shuffle=False, num_workers=n_workers)\n",
        "\n",
        "#model\n",
        "model = PSPNet(layers=50, classes=21).to(device)\n",
        "\n",
        "#loss\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "#optimizer\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
        "n_eps = 30\n",
        "\n",
        "#meter\n",
        "train_loss_meter = AverageMeter()\n",
        "intersection_meter = AverageMeter()\n",
        "union_meter = AverageMeter()\n",
        "target_meter = AverageMeter()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-vf73-nLFrkm"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Training thôi ...</h3>\n",
        "</div>\n",
        "Tham khảo thêm cách viết code training trong PyTorch: https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iyoAzYen1G3p"
      },
      "outputs": [],
      "source": [
        "#training script 2\n",
        "\n",
        "for ep in range(1, 1+n_eps):\n",
        "    train_loss_meter.reset()\n",
        "    intersection_meter.reset()\n",
        "    union_meter.reset()\n",
        "    target_meter.reset()\n",
        "    model.train()\n",
        "\n",
        "    for batch_id, (x, y) in enumerate(tqdm(trainloader), start=1):\n",
        "        #qua trinh hoc mo hinh theo batch\n",
        "        optimizer.zero_grad()\n",
        "        n = x.shape[0]\n",
        "        x = x.to(device).float()\n",
        "        y = y.to(device).long()\n",
        "        y_hat_mask, main_loss, aux_loss = model(x, y)\n",
        "        loss = main_loss + aux_loss*0.4\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        #save metrics\n",
        "        with torch.no_grad():\n",
        "            train_loss_meter.update(loss.item())\n",
        "            intersection, union, target = intersectionAndUnionGPU(y_hat_mask.float(), y.float(), 21)\n",
        "            intersection_meter.update(intersection)\n",
        "            union_meter.update(union)\n",
        "            target_meter.update(target)\n",
        "\n",
        "    #compute iou, dice\n",
        "    with torch.no_grad():\n",
        "        iou_class = intersection_meter.sum / (union_meter.sum + 1e-10) #vector 21D\n",
        "        dice_class = (2 * intersection_meter.sum) / (intersection_meter.sum + union_meter.sum + 1e-10) #vector 21D\n",
        "\n",
        "        mIoU = torch.mean(iou_class) #mean vector 21D\n",
        "        mDice = torch.mean(dice_class) #mean vector 21D\n",
        "\n",
        "    print(\"EP {}, train loss = {} IoU = {}, dice = {}\".format(ep, train_loss_meter.avg, mIoU, mDice))\n",
        "    \n",
        "    if ep >= 25:\n",
        "        torch.save(model.state_dict(), \"modelPSPNet_ep_{}.pth\".format(ep))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Viết code đưa ra kết quả trên tập dữ liệu test</h3>\n",
        "</div>\n",
        "Lưu ý: Trong quá trình test cần chuyển model về chế độ eval và nên đặt torch.no_grad(). Điểm khác nhau giữa model.eval() và torch.no_grad() là model.eval() sẽ đặt chế độ của mô hình về chế độ evaluation (thay đổi các layer như dropout, batch norm, ... khác với quá trình training) còn torch.no_grad() sẽ tắt tính toán đạo hàm của PyTorch để tiết kiệm tính toán. Do vậy khi test mô hình nên chú ý dùng cả 2 câu lệnh này."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model.eval()\n",
        "test_intersection_meter = AverageMeter()\n",
        "test_union_meter = AverageMeter()\n",
        "test_target_meter = AverageMeter()\n",
        "with torch.no_grad():\n",
        "    for batch_id, (x, y) in enumerate(tqdm(testloader), start=1):\n",
        "        n = x.shape[0]\n",
        "        x = x.to(device).float()\n",
        "        y = y.to(device).long()\n",
        "        y_hat = model(x)\n",
        "        y_hat = y_hat.squeeze(1)\n",
        "        y_hat_mask = y_hat.argmax(dim=1)\n",
        "        \n",
        "        intersection, union, target = intersectionAndUnionGPU(y_hat_mask, y, 21)\n",
        "        test_intersection_meter.update(intersection)\n",
        "        test_union_meter.update(union)\n",
        "        test_target_meter.update(target)\n",
        "        \n",
        "    iou_class = test_intersection_meter.sum / (test_union_meter.sum + 1e-10)\n",
        "    dice_class = 2*test_intersection_meter.sum / (test_intersection_meter.sum + test_union_meter.sum + 1e-10)\n",
        "\n",
        "    mIoU = torch.mean(iou_class)\n",
        "    mDice = torch.mean(dice_class)\n",
        "        \n",
        "print(\"TEST: IoU = {}, dice = {}\".format(mIoU, mDice))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OdAA1WJJKGir"
      },
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Viết code hiển thị kết quả dự đoán</h3>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EBIp2CbXKFGH"
      },
      "outputs": [],
      "source": [
        "#predict\n",
        "import random\n",
        "id = random.randint(test_dataset.__len__())\n",
        "with torch.no_grad():\n",
        "    model.eval()\n",
        "    x, y = test_dataset.__getitem__(id)\n",
        "    y_predict = model(x.unsqueeze(0).to(device)).argmax(dim=1).squeeze().cpu().numpy()\n",
        "    for i in np.unique(y_predict).tolist():\n",
        "        print(VOC_CLASSES[i])\n",
        "    color_mask_predict = np.zeros((*y_predict.shape, 3))\n",
        "    for i, color in enumerate(VOC_COLORMAP):\n",
        "        color_mask_predict[y_predict==i] = np.array(color)\n",
        "    color_mask = np.zeros((*y_predict.shape, 3))\n",
        "    for i, color in enumerate(VOC_COLORMAP):\n",
        "        color_mask[y==i] = np.array(color)\n",
        "    plt.subplot(1,3,1)\n",
        "    plt.imshow(unorm(x).permute(1, 2, 0))\n",
        "    plt.subplot(1,3,2)\n",
        "    plt.imshow(color_mask)\n",
        "    plt.subplot(1,3,3)\n",
        "    plt.imshow(color_mask_predict)\n",
        "    plt.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "torch_seg",
      "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.15"
    },
    "vscode": {
      "interpreter": {
        "hash": "ce06b65ffa354aec5ed2625f33c72f9dd46b0d2ee10cd3481aa94d5dfff5c8d3"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}