B2_basic_pipeline_PASCALVOC_ResUNet.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 torchmetrics\n",
        "!pip install segmentation_models_pytorch\n",
        "!pip install albumentations\n",
        "!pip install timm"
      ]
    },
    {
      "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 torchmetrics\n",
        "from torchmetrics import Dice, JaccardIndex\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\n",
        "import timm"
      ]
    },
    {
      "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>Lập trình mô hình UNet với backbone là các mạng ResNet, DenseNet, EfficientNet, ... từ thư viện timm</h3>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3nSSXeMMmXy5"
      },
      "outputs": [],
      "source": [
        "#model UNet\n",
        "def unet_block(in_channels, out_channels):\n",
        "    return nn.Sequential(\n",
        "        nn.Conv2d(in_channels, out_channels, 3, 1, 1),\n",
        "        nn.ReLU(),\n",
        "        nn.Conv2d(out_channels, out_channels, 3, 1, 1),\n",
        "        nn.ReLU()\n",
        "    )\n",
        "\n",
        "class ResUnet(nn.Module):\n",
        "    def __init__(self, n_classes):\n",
        "        super().__init__()\n",
        "        self.n_classes = n_classes\n",
        "        self.backbone = timm.create_model(\"resnet50\", pretrained=True, features_only=True)\n",
        "        self.upsample = nn.Upsample(scale_factor=2, mode=\"bilinear\")\n",
        "        self.block_neck = unet_block(2048, 1024)\n",
        "        self.block_up1 = unet_block(1024+1024, 512)\n",
        "        self.block_up2 = unet_block(512+512, 256)\n",
        "        self.block_up3 = unet_block(256+256, 128)\n",
        "        self.block_up4 = unet_block(128+64, 64)\n",
        "        self.conv_cls = nn.Conv2d(64, self.n_classes, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x1, x2, x3, x4, x5 = self.backbone(x)\n",
        "        #x1 (B, 64, 128, 128) size/2\n",
        "        #x2 (B, 256, 64, 64) size/4\n",
        "        #x3 (B, 512, 32, 32) size/16\n",
        "        #x4 (B, 1024, 16, 16) size/32\n",
        "        #x5 (B, 2048, 8, 8) size/64\n",
        "        x = self.block_neck(x5) # x (B, 1024, 8, 8)\n",
        "        x = torch.cat([x4, self.upsample(x)], dim=1)\n",
        "        x = self.block_up1(x)\n",
        "        x = torch.cat([x3, self.upsample(x)], dim=1)\n",
        "        x = self.block_up2(x)\n",
        "        x = torch.cat([x2, self.upsample(x)], dim=1)\n",
        "        x = self.block_up3(x)\n",
        "        x = torch.cat([x1, self.upsample(x)], dim=1)\n",
        "        x = self.block_up4(x)\n",
        "        x = self.conv_cls(x) #size/2\n",
        "        x = self.upsample(x)\n",
        "        return x\n",
        "\n",
        "# model2 = ResUnet(21)\n",
        "# x = torch.rand(2, 3, 256, 256)\n",
        "# y = model2(x)\n",
        "# print(y.shape)"
      ]
    },
    {
      "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"
      ]
    },
    {
      "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</h3>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9S0zUD210VJ5"
      },
      "outputs": [],
      "source": [
        "#accuracy fn\n",
        "def accuracy_function(preds, targets):\n",
        "    preds_flat = preds.flatten()\n",
        "    targets_flat = targets.flatten()\n",
        "    acc = torch.sum(preds_flat == targets_flat)\n",
        "    return acc/targets_flat.shape[0]"
      ]
    },
    {
      "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\" 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 = ResUnet(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",
        "#metrics\n",
        "dice_fn = torchmetrics.Dice(num_classes=21, average=\"macro\").to(device)\n",
        "iou_fn = torchmetrics.JaccardIndex(num_classes=21, task=\"multiclass\", average=\"macro\").to(device)\n",
        "\n",
        "#meter\n",
        "acc_meter = AverageMeter()\n",
        "train_loss_meter = AverageMeter()\n",
        "dice_meter = AverageMeter()\n",
        "iou_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": [
        "for ep in range(1, 1+n_eps):\n",
        "    acc_meter.reset()\n",
        "    train_loss_meter.reset()\n",
        "    dice_meter.reset()\n",
        "    iou_meter.reset()\n",
        "    model.train()\n",
        "\n",
        "    for batch_id, (x, y) in enumerate(tqdm(trainloader), start=1):\n",
        "        optimizer.zero_grad()\n",
        "        n = x.shape[0]\n",
        "        x = x.to(device).float()\n",
        "        y = y.to(device).long()\n",
        "        y_hat = model(x) #(B, C, H, W)\n",
        "        loss = criterion(y_hat, y) #(B, C, H, W) >< (B, H, W)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            y_hat_mask = y_hat.argmax(dim=1).squeeze() # (B, C, H, W) -> (B, 1, H, W) -> (B, H, W)\n",
        "            dice_score = dice_fn(y_hat_mask, y.long())\n",
        "            iou_score = iou_fn(y_hat_mask, y.long())\n",
        "            accuracy = accuracy_function(y_hat_mask, y.long())\n",
        "\n",
        "            train_loss_meter.update(loss.item(), n)\n",
        "            iou_meter.update(iou_score.item(), n)\n",
        "            dice_meter.update(dice_score.item(), n)\n",
        "            acc_meter.update(accuracy.item(), n)\n",
        "\n",
        "    print(\"EP {}, train loss = {}, accuracy = {}, IoU = {}, dice = {}\".format(\n",
        "        ep, train_loss_meter.avg, acc_meter.avg, iou_meter.avg, dice_meter.avg\n",
        "    ))\n",
        "    if ep >= 25:\n",
        "        torch.save(model.state_dict(), \"modelUNet_ep_{}.pth\".format(ep))"
      ]
    },
    {
      "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": {
      "name": "python",
      "version": "3.8.15"
    },
    "vscode": {
      "interpreter": {
        "hash": "ce06b65ffa354aec5ed2625f33c72f9dd46b0d2ee10cd3481aa94d5dfff5c8d3"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}