transformer_representation_part_2_exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f68eba8",
   "metadata": {},
   "source": [
    "<h1 style=\"font-family: serif\">Exercise: Transformer Representation</h1>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "587a7c24",
   "metadata": {},
   "source": [
    "In this exercise, we'll learn how to extract features out of a pre-trained transformer model and use those features for downstream tasks. For this exercise, we'll use [*DINOv2: A Self-supervised Vision Transformer Model*](https://dinov2.metademolab.com/) by *meta*. This model is trained in a teacher-student paradigm, without any supervision, and it produces features suitable for different downstream tasks like image classification, depth estimation, semantic segmentation, etc.\n",
    "<br><br>**Note:** DINOv2 makes 14x14 patches out of an input image, and then produce features for each patch (not for each pixel)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cc4edd5",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-danger\">\n",
    "    Please switch to the <code>03_learned_representations</code> environment.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfc8e70f",
   "metadata": {},
   "source": [
    "<h2 style=\"font-family: serif\">II. Part Two</h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f54ae58",
   "metadata": {},
   "source": [
    "In the second part, we will train a model using the DINOv2 extracted features as inputs. The task will be semantic segmentation over the input image patches.\n",
    "<br>For model evaluation, we are using metrics from [**Segmentation Models**](https://smp.readthedocs.io/en/latest/metrics.html) package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e38ce30",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision.transforms import v2 as tv_transforms2\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import utils\n",
    "import metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60da2527",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# to have interactive plots\n",
    "%matplotlib widget\n",
    "plt.ioff()\n",
    "\n",
    "SEED = 2024\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "333aead9",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# color map for visualization of the ground truth masks\n",
    "cm, colors = utils.get_colormap()\n",
    "cm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0dccb2c",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "We are using data provided by [*Dense cellular segmentation for EM using 2D-3D neural network ensembles*](https://leapmanlab.github.io/dense-cell/).\n",
    "<br>The data contains *tiff* files in train and evaluation sets along with their ground truth masks. Masks include dense annotations for seven classes.\n",
    "<br>Images have a resolution of 800 x 800 pixels which are a bit large to fit in memory or GPU 😁 . However, we'll transform them into a smaller manageable resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9bb69fc",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# dataset class definition\n",
    "\n",
    "class TiffDataset(Dataset):\n",
    "    def __init__(self, image_path, label_path, input_size, train=True):\n",
    "        self.images = utils.get_images_from_tiff(image_path, to_rgb=True)  # numpy array, channel last\n",
    "        self.gt_masks = utils.get_images_from_tiff(label_path, to_rgb=False)\n",
    "        self.input_size = input_size\n",
    "        self.train = train\n",
    "        self.img_h = self.images.shape[1]\n",
    "        self.img_w = self.images.shape[2]\n",
    "        self.mean = None\n",
    "        self.std = None\n",
    "        self.mean, self.std = self.get_mean_std()\n",
    "        self.base_transform = tv_transforms2.Compose([\n",
    "            tv_transforms2.ToImage(),\n",
    "        ])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        image = self.images[index]\n",
    "        mask = self.gt_masks[index]\n",
    "\n",
    "        return self.apply_transform(image, mask)\n",
    "\n",
    "    def apply_transform(self, image, mask):\n",
    "        # check channel dimension\n",
    "        if len(image.shape) < 3:\n",
    "            image = np.expand_dims(image, axis=-1)\n",
    "        if len(mask.shape) < 3:\n",
    "            mask = np.expand_dims(mask, axis=-1)\n",
    "\n",
    "        image, mask = self.base_transform(image, mask)\n",
    "        # resizing\n",
    "        image = tv_transforms2.functional.resize(image, size=self.input_size)\n",
    "        mask = tv_transforms2.functional.resize(\n",
    "            mask, size=self.input_size,\n",
    "            interpolation=tv_transforms2.InterpolationMode.NEAREST_EXACT, antialias=False\n",
    "        )\n",
    "        # to tensor\n",
    "        image = tv_transforms2.functional.to_dtype(image, dtype=torch.float32, scale=True)\n",
    "        mask = tv_transforms2.functional.to_dtype(mask, dtype=torch.long, scale=False)\n",
    "        # normalizing\n",
    "        image = tv_transforms2.functional.normalize(image, self.mean, self.std)\n",
    "        assert torch.isin(torch.unique(mask), torch.arange(7)).all()\n",
    "\n",
    "        return image, mask.squeeze(0)\n",
    "\n",
    "    def get_mean_std(self):\n",
    "        _min = self.images.min()\n",
    "        _max = self.images.max()\n",
    "        scaled_imgs = (self.images - _min) / (_max - _min)\n",
    "    \n",
    "        return scaled_imgs.mean(), scaled_imgs.std()\n",
    "\n",
    "    def get_class_weights(self):\n",
    "        _, weights = np.unique(self.gt_masks, return_counts=True)\n",
    "        weights = 1 / weights\n",
    "        # normalize weights\n",
    "        weights = weights / weights.sum()\n",
    "\n",
    "        return weights\n",
    "\n",
    "    def plot_sample(self, cm=None):\n",
    "        image, mask = self.__getitem__(0)\n",
    "        cm = cm or \"Dark2\"\n",
    "        fig, axes = plt.subplots(1, 3, figsize=(12.5, 5.0), layout=\"compressed\")\n",
    "        fig.canvas.toolbar_position = \"right\"\n",
    "        fig.canvas.header_visible = False\n",
    "        fig.canvas.footer_visible = False\n",
    "        axes[0].imshow(image[0], cmap=\"grey\")\n",
    "        axes[0].set_title(\"Image\")\n",
    "        axes[1].imshow(mask, cmap=cm, vmax=7, interpolation=\"none\")\n",
    "        axes[1].set_title(\"Label\")\n",
    "        axes[2].imshow(image[0], interpolation=\"none\", cmap=\"grey\")\n",
    "        axes[2].imshow(mask, alpha=0.45, cmap=cm, vmax=7, interpolation=\"none\")\n",
    "        axes[2].set_title(\"Overlay\")\n",
    "\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17e36257",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ground truth classes and their labels\n",
    "num_classes = 7\n",
    "classes = {\n",
    "    \"background\": 0,\n",
    "    \"cell\": 1,\n",
    "    \"mitochondrion\": 2,\n",
    "    \"alpha granule\": 3,\n",
    "    \"canalicular channel\": 4,\n",
    "    \"dense granule\": 5,\n",
    "    \"dense granule core\": 6\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ac9bf29",
   "metadata": {},
   "source": [
    "## Load the Pre-trained Transformer Model\n",
    "We use pre-trained DINOv2 small model for feature extraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f4c6573",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68175b14",
   "metadata": {},
   "outputs": [],
   "source": [
    "dinov2 = torch.hub.load(\"facebookresearch/dinov2\", \"dinov2_vits14_reg\").to(device)\n",
    "dinov2.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "defbcdc5",
   "metadata": {},
   "source": [
    "We will use the dino's `get_intermediate_layers` method to extract features from the DINOv2 model.  \n",
    "- The first parameter is an input image batch. \n",
    "- The second parameter, `n`, points to model's layer(s) to extract features from (layers or n last layers to take).  \n",
    "- If `reshape=True`, the features will be returned as a batch of 3D : (F-size, W, H), else it will be 2D ((W x H), F-size).  \n",
    "- We don't want the class token, so `return_class_token=False`.  \n",
    "<br><br>\n",
    "This method returns a tuple of features with each element points to a requested layer.\n",
    "<br> See the code [*here*](https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L298)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99ec8254",
   "metadata": {},
   "outputs": [],
   "source": [
    "help(dinov2.get_intermediate_layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dbc3c15",
   "metadata": {},
   "source": [
    "## Segmentation Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7efaf3d9",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "  <h3>Task 1.1: Implementing a Model for Segmentation</h3>\n",
    "  <p>\n",
    "      It's time to define our segmentation model!\n",
    "      <br>Start with a simple model, use convolution layers, and remember that the input has a resolution of (num_patches × num_patches).\n",
    "  </p>\n",
    "<p><i>\n",
    "   Please refer to the <a href=\"https://pytorch.org/docs/stable/nn.html\"><b>Pytorch</b> documentation</a>\n",
    "</i>.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c958a63b",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# the base class for your model to derive from (just gives you the number of trainable params) :)\n",
    "class BaseNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def count_parameters(self, trainable=True):\n",
    "        params = [\n",
    "            param.numel()\n",
    "            for param in self.parameters() if param.requires_grad == trainable\n",
    "        ]\n",
    "        return sum(params), params\n",
    "\n",
    "    def __repr__(self):\n",
    "        params = self.count_parameters()\n",
    "        return f\"{super().__repr__()}\\ntrainable params: {params[0]:,d} {params[1]}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a2090c7",
   "metadata": {
    "lines_to_next_cell": 1,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# complete the model implementation\n",
    "\n",
    "# class Net(BaseNet):\n",
    "#     def __init__(self, ...):\n",
    "#         super().__init__()\n",
    "#         self.input_dim = num_patches\n",
    "#         self.in_channels = in_channels\n",
    "#         self.n_classes = n_classes\n",
    "#         \n",
    "#         self.conv1 = nn.Sequential(\n",
    "#              nn.Conv2d(\n",
    "#                   self.in_channels, 256, kernel_size=3, padding=1, bias=False\n",
    "#               ),\n",
    "#              nn.BatchNorm2d(256),\n",
    "#              nn.LeakyReLU(negative_slope=0.01),\n",
    "#         )\n",
    "#        # add a similar module as above. Note that input channels to this module will be the same as what is output channels\n",
    "#        self.conv2 = ...\n",
    "#        # add a similar module. Note that output channels of this should be same as input channels of self.conv_out\n",
    "#        self.conv3 = ...\n",
    "#        # segmentation output should have channels equal to the number of classes.        \n",
    "#        self.conv_out = nn.Conv2d(\n",
    "#            64, self.n_classes, kernel_size=1, bias=False\n",
    "#        )\n",
    "#\n",
    "#    def forward(self, x):\n",
    "#        # input will be a tensor of b x (num_patches^2) x 384 dims.\n",
    "#        x = x.reshape(-1, self.input_dim, self.input_dim, self.in_channels)\n",
    "#        x = x.permute(0, 3, 1, 2)\n",
    "#        out = self.conv1(x)\n",
    "#        out = self.conv2(out)\n",
    "#        out = self.conv3(out)\n",
    "#        out = self.conv_out(out)\n",
    "\n",
    "#        return out\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae2050e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# DINOv2 trained on image patches of size 14 x 14. Therefore, the input image size should be divisible by 14. \n",
    "# dinov2_vits14_reg specs:\n",
    "feature_dim = 384\n",
    "patch_size = 14\n",
    "# to reduce original image resolution to integer number of patches\n",
    "num_patches = 30\n",
    "\n",
    "input_size = patch_size * num_patches\n",
    "print(f\"Dino input image size: {input_size}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "134067eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Net(num_patches, feature_dim, num_classes).to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2989381",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"torch.cuda.memory_allocated: %fGB\"%(torch.cuda.memory_allocated(0)/1024/1024/1024))\n",
    "print(\"torch.cuda.memory_reserved: %fGB\"%(torch.cuda.memory_reserved(0)/1024/1024/1024))\n",
    "print(\"torch.cuda.max_memory_reserved: %fGB\"%(torch.cuda.max_memory_reserved(0)/1024/1024/1024))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00edb613",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa6eb88d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the data resides in this path: \"/group/dl4miacourse/platelet_data\"\n",
    "# train dataset\n",
    "train_dataset = TiffDataset(\n",
    "    image_path=\"/group/dl4miacourse/platelet_data/train-images.tif\",\n",
    "    label_path=\"/group/dl4miacourse/platelet_data/train-labels.tif\",\n",
    "    input_size=input_size\n",
    ")\n",
    "\n",
    "print(f\"number of images: {len(train_dataset)}\")\n",
    "train_dataset.plot_sample(cm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f42f39e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyper-params\n",
    "batch_size = 16\n",
    "lr = 1e-3\n",
    "epochs = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cbac5f8",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "  <h3>Task 1.2: Training</h3>\n",
    "  <p>\n",
    "      The training loop is here already! You just need to setup the optimizer and the loss function.\n",
    "      <br><b>Note:</b> In segmentation tasks, usually some classes including much more pixels than others e.g. background. So, we need to take care of this class imbalances giving weights to each classes. To do this use <code>train_dataset.get_class_weights()</code> to get the class's weights.\n",
    "  </p>\n",
    "<p><i>\n",
    "   Please refer to the pytorch documentations: <a href=\"https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss\">CrossEntropyLoss</a>, <a href=\"https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html#torch.optim.AdamW\">AdamW</a> \n",
    "</i>.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876e61e7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# insert your code here\n",
    "# optim = ...  # you can use Adam or AdamW\n",
    "\n",
    "# class_weights = torch.from_numpy(train_dataset.get_class_weights()).to(torch.float32).to(device)\n",
    "# loss_fn = ...  # use CrossEntropyLoss with the weight param equals to the class weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3961b571",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset, batch_size, shuffle=True)\n",
    "num_batches = len(train_loader)\n",
    "print(f\"number of batches: {num_batches}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6dcd980",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "To visualize the training progress, we use an interactive ploting method: \n",
    "<br> at each plot update, we remove the previous line and draw the new one, \n",
    "then we update the plot by calling `fig_loss.canvas.draw()` and `fig_loss.canvas.flush_events()` ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3918e71b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# preparing training loss plot\n",
    "fig_loss, ax_loss = plt.subplots(1, 1, figsize=(11, 6), layout=\"compressed\")\n",
    "fig_loss.canvas.toolbar_position = \"right\"\n",
    "fig_loss.canvas.header_visible = False\n",
    "fig_loss.canvas.footer_visible = False\n",
    "ax_loss.set_title(\"Training Loss\")\n",
    "ax_loss.set_xlabel(\"Iterations\")\n",
    "ax_loss.set_ylabel(\"Loss\")\n",
    "ax_loss.grid(alpha=0.3)\n",
    "loss_line = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a8dcab",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.show()\n",
    "losses = []\n",
    "tps, fps, fns, tns = [], [], [], []\n",
    "\n",
    "for e in tqdm(range(epochs), desc=\"Training Epochs\"):\n",
    "    for batch_idx, (image, gt_masks) in enumerate(train_loader):\n",
    "        image = image.to(device)\n",
    "        gt_masks = gt_masks.to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            features = dinov2.get_intermediate_layers(image, 1, return_class_token=False)[0]\n",
    "        out = model(features)\n",
    "\n",
    "        # here we scale up the model output into the gt_mask size\n",
    "        out_upscaled = F.interpolate(out, size=input_size, mode=\"bilinear\", align_corners=False)\n",
    "\n",
    "        loss = loss_fn(out_upscaled, gt_masks)\n",
    "        optim.zero_grad()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "\n",
    "        # for metrics\n",
    "        tp, fp, fn, tn = metrics.get_stats(\n",
    "            out_upscaled.argmax(dim=1),\n",
    "            gt_masks, mode=\"multiclass\",\n",
    "            num_classes=num_classes\n",
    "        )\n",
    "        tps.append(tp)\n",
    "        fps.append(fp)\n",
    "        fns.append(fn)\n",
    "        tns.append(tn)\n",
    "\n",
    "        losses.append(loss.item())\n",
    "        if batch_idx == 0 or batch_idx % 5 == 0:\n",
    "            print(f\"loss: {loss.item():.5f}\", end=\"\\r\")\n",
    "            if loss_line:\n",
    "                loss_line[0].remove()\n",
    "            loss_line = ax_loss.plot(losses, color=\"dodgerblue\", label=\"Train Loss\")\n",
    "            fig_loss.canvas.draw()\n",
    "            fig_loss.canvas.flush_events()\n",
    "            ax_loss.legend()\n",
    "    # end of one epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fd50929",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# calculate metrics\n",
    "accs = []\n",
    "f1_scores = []\n",
    "ious = []\n",
    "\n",
    "for i in range(0, len(tps), num_batches):\n",
    "    epoch_tp = torch.concat(tps[i: i + num_batches])\n",
    "    epoch_fp = torch.concat(fps[i: i + num_batches])\n",
    "    epoch_fn = torch.concat(fns[i: i + num_batches])\n",
    "    epoch_tn = torch.concat(tns[i: i + num_batches])\n",
    "    accs.append(\n",
    "        metrics.accuracy(epoch_tp, epoch_fp, epoch_fn, epoch_tn, reduction=None).mean(dim=0).numpy()\n",
    "    )\n",
    "    f1_scores.append(\n",
    "        metrics.f1_score(epoch_tp, epoch_fp, epoch_fn, epoch_tn, reduction=None).mean(dim=0).numpy()\n",
    "    )\n",
    "    ious.append(\n",
    "        metrics.iou_score(epoch_tp, epoch_fp, epoch_fn, epoch_tn, reduction=None).mean(dim=0).numpy()\n",
    "    )\n",
    "\n",
    "accs = np.vstack(accs)\n",
    "f1_scores = np.vstack(f1_scores)\n",
    "ious = np.vstack(ious)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "437fd40f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot metrics\n",
    "def plot_metric(metric, ax, title, colors, class_labels):\n",
    "    # metric shape: num_epochs x num_classes\n",
    "    num_classes = metric.shape[1]\n",
    "    for c in range(num_classes):\n",
    "        ax.plot(metric[:, c] * 100, color=colors[c], label=class_labels[c], lw=1.2)\n",
    "    # epoch average\n",
    "    ax.plot(metric.mean(axis=1) * 100, color=\"maroon\", label=\"Average\", lw=1.5)\n",
    "\n",
    "    ax.set_title(title, fontweight=\"bold\")\n",
    "    ax.set_xlabel(\"Epochs\")\n",
    "    ax.set_ylabel(\"%\")\n",
    "    ax.grid(alpha=0.3)\n",
    "    ax.legend(loc=\"lower right\", fontsize=8)\n",
    "\n",
    "fig_metrics, axes = plt.subplots(1, 3, figsize=(16, 5), layout=\"compressed\")\n",
    "fig_metrics.canvas.toolbar_position = \"right\"\n",
    "fig_metrics.canvas.header_visible = False\n",
    "fig_metrics.canvas.footer_visible = False\n",
    "\n",
    "plot_metric(accs, axes[0], \"Accuracy\", colors, list(classes.keys()))\n",
    "plot_metric(f1_scores, axes[1], \"F1 Score\", colors, list(classes.keys()))\n",
    "plot_metric(ious, axes[2], \"IoU\", colors, list(classes.keys()))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c53fcb2",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "  <h3>Checkpoint 1</h3>\n",
    "  <p>Congratulations! You managed to train a segmentation model using the DINOv2 features as inputs.<br>Let's evaluate your model! 😁</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee1809d8",
   "metadata": {},
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4253706a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test dataset \n",
    "# Note: we are using normalization stats from the train dataset.\n",
    "mean = train_dataset.mean\n",
    "std = train_dataset.std\n",
    "\n",
    "# the data resides in this path: \"/group/dl4miacourse/platelet_data\"\n",
    "test_dataset = TiffDataset(\n",
    "    image_path=\"/group/dl4miacourse/platelet_data/eval-images.tif\",\n",
    "    label_path=\"/group/dl4miacourse/platelet_data/eval-labels.tif\",\n",
    "    input_size=input_size, train=False\n",
    ")\n",
    "test_dataset.mean = mean\n",
    "test_dataset.std = std\n",
    "\n",
    "test_loader = DataLoader(test_dataset, batch_size, shuffle=False)\n",
    "\n",
    "print(f\"number of images: {len(test_dataset)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91ab90bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = []\n",
    "tps, fps, fns, tns = [], [], [], []\n",
    "\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7082de0a",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "  <h3>Task 1.3: Model Evaluation</h3>\n",
    "  <p>\n",
    "      Given the train loop, this is a very easy one!\n",
    "      <br>You need to pass the input image to the DINOv2 to get the features, and then pass the features to your model. Done!\n",
    "  </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4097eabd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# complete the testing code\n",
    "# for image, gt_masks in test_loader:\n",
    "#    image = image.to(device)\n",
    "#    gt_masks = gt_masks.to(device)\n",
    "\n",
    "#    with torch.no_grad():\n",
    "#        # pass image to the DINO to get the features\n",
    "#        features = ...\n",
    "#        # pass features to your model\n",
    "#        out = ...\n",
    "#\n",
    "#    out_upscaled = F.interpolate(out, size=input_size, mode=\"bilinear\", align_corners=False)\n",
    "#    loss = loss_fn(out_upscaled, gt_masks)\n",
    "#    Don't forget your metrics!\n",
    "#    tp, fp, fn, tn = metrics.get_stats(\n",
    "#        out_upscaled.argmax(dim=1),\n",
    "#        gt_masks, mode=\"multiclass\",\n",
    "#        num_classes=num_classes\n",
    "#    )\n",
    "#    tps.append(tp)\n",
    "#    fps.append(fp)\n",
    "#    fns.append(fn)\n",
    "#    tns.append(tn)\n",
    "#    losses.append(loss.item())\n",
    "#\n",
    "# print(f\"Evaluation average loss: {np.mean(losses):.5f}\", end=\"\\r\")\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89cac04d",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# calculate metrics\n",
    "tps = torch.concat(tps)\n",
    "fps = torch.concat(fps)\n",
    "fns = torch.concat(fns)\n",
    "tns = torch.concat(tns)\n",
    "\n",
    "accs = metrics.accuracy(tps, fps, fns, tns, reduction=None).mean(dim=0).numpy()\n",
    "f1_scores = metrics.f1_score(tps, fps, fns, tns, reduction=None).mean(dim=0).numpy()\n",
    "ious = metrics.iou_score(tps, fps, fns, tns, reduction=None).mean(dim=0).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0cacdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot evaluation metrics\n",
    "def plot_evaluation(metric, ax, title, colors, class_labels):\n",
    "    bp = ax.bar(class_labels, height=metric * 100, color=colors, width=0.65)\n",
    "    ax.bar_label(bp, label_type=\"edge\", fmt=\"%.2f\")\n",
    "    ax.set_title(title, fontweight=\"bold\")\n",
    "    ax.tick_params(axis=\"x\", labelrotation=45)\n",
    "\n",
    "\n",
    "fig_metrics, axes = plt.subplots(1, 3, figsize=(16, 5), layout=\"compressed\")\n",
    "fig_metrics.canvas.toolbar_position = \"right\"\n",
    "fig_metrics.canvas.header_visible = False\n",
    "fig_metrics.canvas.footer_visible = False\n",
    "\n",
    "plot_evaluation(accs, axes[0], \"Accuracy\", colors, list(classes.keys()))\n",
    "plot_evaluation(f1_scores, axes[1], \"F1 Score\", colors, list(classes.keys()))\n",
    "plot_evaluation(ious, axes[2], \"IoU\", colors, list(classes.keys()))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "369a6392",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot a sample of model's segmentation result\n",
    "img, gt = test_dataset.__getitem__(0)\n",
    "img = img.unsqueeze(0).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    features = dinov2.get_intermediate_layers(img, 1, return_class_token=False)[0]\n",
    "    out = model(features)\n",
    "out_upscaled = F.interpolate(out, size=input_size, mode=\"bilinear\", align_corners=False)\n",
    "pred = out_upscaled.squeeze(0)\n",
    "pred = pred.argmax(dim=0)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(11, 6), layout=\"compressed\")\n",
    "fig.canvas.header_visible = False\n",
    "axes[0].imshow(img[0, 0].cpu(), cmap=\"grey\", interpolation=\"none\")\n",
    "axes[0].imshow(pred.cpu(), cmap=cm, interpolation=\"none\", vmax=num_classes - 1, alpha=0.5)\n",
    "axes[0].set_title(\"Prediction\")\n",
    "axes[1].imshow(img[0, 0].cpu(), cmap=\"grey\", interpolation=\"none\")\n",
    "axes[1].imshow(gt, cmap=cm, interpolation=\"none\", vmax=num_classes - 1, alpha=0.5)\n",
    "axes[1].set_title(\"Ground Truth\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3823133",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "  <h3>Checkpoint 2</h3>\n",
    "  <p>Congratulations × 2 !! <br>Now we learned that DINOv2 or in general vision transformer models, can extract meaningful features out of our dataset images, even though they are usually trained on natural images.\n",
    "  <br>We can use those features for many downstream tasks, including semantic segmentation. \n",
    "  </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "578729c8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "tags,-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all",
   "text_representation": {
    "extension": ".py",
    "format_name": "percent"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}