B6_PAN.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "!pip install segmentation_models_pytorch\n",
    "!pip install albumentations"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Import thư viện</h3>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "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": {},
   "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",
    "\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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "trainsize = 256\n",
    "\n",
    "train_transform = A.Compose([\n",
    "    A.Resize(width=trainsize, height=trainsize),\n",
    "    A.ShiftScaleRotate(),\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",
    "])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "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))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "train_dataset = PascalVOCSearchDataset(image_set=\"train\", download=False, 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()\n",
    "print(mask.unique())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "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": {},
   "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": {},
   "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 "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "#load data\n",
    "batch_size = 8\n",
    "#batch_size = 4\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 = smp.create_model(\"PAN\", \"timm-efficientnet-b4\", \"imagenet\", 3, 21).to(device)\n",
    "\n",
    "#loss\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "# criterion = smp.losses.DiceLoss(mode=\"multiclass\", classes=21) #diceloss = 1-dicescore\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()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "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",
    "        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",
    "        #save metrics\n",
    "        with torch.no_grad():\n",
    "            train_loss_meter.update(loss.item())\n",
    "            y_hat_mask = y_hat.argmax(dim=1).squeeze(1) # (B, C, H, W) -> (B, 1, H, W) -> (B, H, W)\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(), \"modelDeepLab_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))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "#predict\n",
    "import random\n",
    "id = random.randint(0, 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()\n",
    "    intersection, union, target = intersectionAndUnionGPU(y_predict.float(), y.to(device).float(), 21)\n",
    "    iou_class = intersection / (union + 1e-10)\n",
    "    dice_class = 2*intersection / (intersection + union + 1e-10)\n",
    "    mIoU = torch.mean(iou_class)\n",
    "    mDice = torch.mean(dice_class)\n",
    "    print(\"IoU = {}\\n dice = {}\".format(iou_class, dice_class))\n",
    "    y_predict = y_predict.cpu().numpy()\n",
    "#     print(np.unique(y_predict))\n",
    "    for i in np.unique(y_predict):\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": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 2
}