B2_basic_pipeline.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Kiểm tra GPU được cấp phát</h3>\n",
        "</div>"
      ],
      "metadata": {
        "id": "BQNLZ86t9R3G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "GOVtzLUjjZ1W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "vigGQG1--GdX"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tChf3YMRYr4m"
      },
      "outputs": [],
      "source": [
        "!pip install torchmetrics\n",
        "!pip install segmentation_models_pytorch\n",
        "!pip install albumentations"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Download dữ liệu</h3>\n",
        "</div>"
      ],
      "metadata": {
        "id": "I_K_s-Yk-nKg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://thor.robots.ox.ac.uk/~vgg/data/pets/images.tar.gz\n",
        "!wget https://thor.robots.ox.ac.uk/~vgg/data/pets/annotations.tar.gz"
      ],
      "metadata": {
        "id": "mQmfCxYGZaTY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3></h3>\n",
        "</div>\n"
      ],
      "metadata": {
        "id": "7dnqP5oJ-yEs"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Giải nén dữ liệu</h3>\n",
        "</div>"
      ],
      "metadata": {
        "id": "ONGrqXmFCAuB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!tar -xf annotations.tar.gz\n",
        "!tar -xf images.tar.gz"
      ],
      "metadata": {
        "id": "fZ5sutIwZnkA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Import thư viện</h3>\n",
        "</div>"
      ],
      "metadata": {
        "id": "n2ZRXglT_5Rm"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "qCnMyzJgZqJr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# mask_path = \"/content/annotations/trimaps/Abyssinian_1.png\"\n",
        "# mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)\n",
        "# plt.imshow(mask)\n",
        "# print(np.unique(mask))"
      ],
      "metadata": {
        "id": "BflMgDcZbPIX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# image_path = \"/content/images/Abyssinian_10.jpg\"\n",
        "# image = cv2.imread(image_path)\n",
        "# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "# plt.imshow(image)\n",
        "# print(image.shape)"
      ],
      "metadata": {
        "id": "Ocq9MrPtcA5G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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ể"
      ],
      "metadata": {
        "id": "Xz7v36fKCFqV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class DogCatDataset(Dataset):\n",
        "    def __init__(self, root_dir, txt_file, transform=None): #transform: augmentation + norm + np.array -> torch.tensor\n",
        "        super().__init__()\n",
        "        self.root_dir = root_dir\n",
        "        self.txt_file = txt_file\n",
        "        self.transform = transform\n",
        "        self.img_path_lst = []\n",
        "        with open(self.txt_file) as file_in:\n",
        "            for line in file_in:\n",
        "                self.img_path_lst.append(line.split(\" \")[0])\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.img_path_lst)\n",
        "    \n",
        "    def __getitem__(self, idx):\n",
        "        image_path = os.path.join(self.root_dir, \"images\", \"{}.jpg\".format(self.img_path_lst[idx]))\n",
        "        mask_path = os.path.join(self.root_dir, \"annotations\", \"trimaps\", \"{}.png\".format(self.img_path_lst[idx]))\n",
        "        image = cv2.imread(image_path)\n",
        "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)\n",
        "        # foreground -> 1\n",
        "        # background 2 -> 0\n",
        "        # 3 -> 1\n",
        "        mask[mask == 2] = 0\n",
        "        mask[mask == 3] = 1\n",
        "        # image (RGB), mask (2D matrix)\n",
        "        if self.transform is not None:\n",
        "            transformed = self.transform(image=image, mask=mask)\n",
        "            transformed_image = transformed['image']\n",
        "            transformed_mask = transformed['mask']\n",
        "        return transformed_image, transformed_mask"
      ],
      "metadata": {
        "id": "GNv3AOficqQK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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/"
      ],
      "metadata": {
        "id": "en9W5v_xChlD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "trainsize = 384\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",
        "])"
      ],
      "metadata": {
        "id": "e4RPW6nMgx4a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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>"
      ],
      "metadata": {
        "id": "r81wtP0PCzUD"
      }
    },
    {
      "cell_type": "code",
      "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))"
      ],
      "metadata": {
        "id": "RSdnanKkjlD3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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>"
      ],
      "metadata": {
        "id": "x_CbSVe6C_bo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset = DogCatDataset(\"/content\", \"/content/annotations/trainval.txt\", train_transform)\n",
        "test_dataset = DogCatDataset(\"/content\", \"/content/annotations/test.txt\", 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()"
      ],
      "metadata": {
        "id": "06Kt7pRTfD8H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Lập trình mô hình UNet cơ bản</h3>\n",
        "</div>"
      ],
      "metadata": {
        "id": "XNubRiEKDHg0"
      }
    },
    {
      "cell_type": "code",
      "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 UNet(nn.Module):\n",
        "    def __init__(self, n_classes):\n",
        "        super().__init__()\n",
        "        self.n_classes = n_classes\n",
        "        self.downsample = nn.MaxPool2d(2)\n",
        "        self.upsample = nn.Upsample(scale_factor=2, mode=\"bilinear\")\n",
        "        self.block_down1 = unet_block(3, 64)\n",
        "        self.block_down2 = unet_block(64, 128)\n",
        "        self.block_down3 = unet_block(128, 256)\n",
        "        self.block_down4 = unet_block(256, 512)\n",
        "        self.block_neck = unet_block(512, 1024)\n",
        "        self.block_up1 = unet_block(1024+512, 512)\n",
        "        self.block_up2 = unet_block(256+512, 256)\n",
        "        self.block_up3 = unet_block(128+256, 128)\n",
        "        self.block_up4 = unet_block(128+64, 64)\n",
        "        self.conv_cls = nn.Conv2d(64, self.n_classes, 1) # -> (B, n_class, H, W)\n",
        "    \n",
        "    def forward(self, x):\n",
        "        # (B, C, H, W)\n",
        "        x1 = self.block_down1(x)\n",
        "        x = self.downsample(x1)\n",
        "        x2 = self.block_down2(x)\n",
        "        x = self.downsample(x2)\n",
        "        x3 = self.block_down3(x)\n",
        "        x = self.downsample(x3)\n",
        "        x4 = self.block_down4(x)\n",
        "        x = self.downsample(x4)\n",
        "\n",
        "        x = self.block_neck(x)\n",
        "\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",
        "\n",
        "        x = self.conv_cls(x)\n",
        "        return x\n",
        "\n",
        "# model = UNet(1)\n",
        "# x = torch.rand(4, 3, trainsize, trainsize)\n",
        "# print(\"Input shape =\", x.shape)\n",
        "# y = model(x).squeeze()\n",
        "# print(\"Output shape = \", y.shape)\n",
        "# # y true (4, 384, 384)\n",
        "# # y hat.squeeze() (4, 1, 384, 384) -> (4, 384, 384)"
      ],
      "metadata": {
        "id": "3nSSXeMMmXy5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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"
      ],
      "metadata": {
        "id": "-OZbtvpnD1ho"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "OJ0zgMSXy-7Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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>"
      ],
      "metadata": {
        "id": "s8u9IBDlEQyF"
      }
    },
    {
      "cell_type": "code",
      "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]"
      ],
      "metadata": {
        "id": "9S0zUD210VJ5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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"
      ],
      "metadata": {
        "id": "MxkzTm38EXCo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#device\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "#load data\n",
        "batch_size = 8\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 = UNet(1).to(device)\n",
        "\n",
        "#loss\n",
        "criterion = nn.BCEWithLogitsLoss()\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=2, average=\"macro\").to(device)\n",
        "iou_fn = torchmetrics.JaccardIndex(num_classes=2, task=\"binary\", average=\"macro\").to(device)\n",
        "\n",
        "#meter\n",
        "acc_meter = AverageMeter()\n",
        "train_loss_meter = AverageMeter()\n",
        "dice_meter = AverageMeter()\n",
        "iou_meter = AverageMeter()"
      ],
      "metadata": {
        "id": "xldeI02fubce"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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"
      ],
      "metadata": {
        "id": "-vf73-nLFrkm"
      }
    },
    {
      "cell_type": "code",
      "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).float()\n",
        "        y_hat = model(x)\n",
        "        y_hat = y_hat.squeeze() # -> logit (-vc, +vc)\n",
        "        loss = criterion(y_hat, y)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            y_hat_mask = y_hat.sigmoid().round().long() # -> mask (0, 1)\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(), \"/content/model_ep_{}.pth\".format(ep))"
      ],
      "metadata": {
        "id": "iyoAzYen1G3p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "Xlx3dVD9I_K8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.eval()\n",
        "test_iou_meter = AverageMeter()\n",
        "test_dice_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).float()\n",
        "        y_hat = model(x)\n",
        "        y_hat = y_hat.squeeze()\n",
        "        y_hat_mask = y_hat.sigmoid().round().long()\n",
        "        dice_score = dice_fn(y_hat_mask, y.long())\n",
        "        iou_score = iou_fn(y_hat_mask, y.long())\n",
        "        test_dice_meter.update(dice_score.item(), n)\n",
        "        test_iou_meter.update(iou_score.item(), n)\n",
        "print(\"TEST: IoU = {}, dice = {}\".format(test_iou_meter.avg, test_dice_meter.avg))"
      ],
      "metadata": {
        "id": "WY4wUsAb5CZQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div class=\"markdown-google-sans\">\n",
        "  <h3>Viết code hiển thị kết quả dự đoán</h3>\n",
        "</div>"
      ],
      "metadata": {
        "id": "OdAA1WJJKGir"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "\n",
        "model.eval()\n",
        "idx = random.randint(0, 100)\n",
        "with torch.no_grad():\n",
        "    x, y = test_dataset[idx]\n",
        "    # print(x.shape, y.shape) (C, H, W) -> (1, C, H, W) -> model\n",
        "    x = x.to(device).float().unsqueeze(0)\n",
        "    y_hat = model(x).squeeze() #(1, 1, H, W) -> (H, W)\n",
        "    y_hat_mask = y_hat.sigmoid().round().long()\n",
        "    # x, y, y_hat_mask\n",
        "    plt.subplot(1, 3, 1)\n",
        "    plt.imshow(unorm(x.squeeze().cpu()).permute(1, 2, 0)) # x (GPU) -> x(CPU)\n",
        "    plt.subplot(1, 3, 2)\n",
        "    plt.imshow(y)\n",
        "    plt.subplot(1, 3, 3)\n",
        "    plt.imshow(y_hat_mask.cpu())"
      ],
      "metadata": {
        "id": "EBIp2CbXKFGH"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}