HOW_TO_USE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1 align=\"center\"> Channel Adaptive Vision Transformer: How to Use </h1>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is a step-by-step guide on how to use the Channel Adaptive Vision Transformer (ChAdaViT) model for image classification. The ChAdaViT model is a vision transformer that can adaptively take as input images from different number of channels, and project them into the same embedding space. This is particularly useful when working with multi-channel images, such as medical microscopy or even geopspatial images with multiple modalities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import hashlib\n",
    "\n",
    "from src.backbones.vit.chada_vit import ChAdaViT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Using device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download weights\n",
    "You can download the model weights under this URL: https://drive.google.com/file/d/1SUfUwerHJlf0vo9mdgM0mRn9TNZkaqXl/view?usp=drive_link   \n",
    "Make sure to download it on the same directory as this notebook, and give the right permissions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter the path of the weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "CKPT_PATH = \"weights.ckpt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can check the hash of the downloaded file here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_hash(file_path, expected_hash):\n",
    "    md5 = hashlib.md5()\n",
    "    with open(file_path, \"rb\") as f:\n",
    "        while chunk := f.read(4096):\n",
    "            md5.update(chunk)\n",
    "    return md5.hexdigest() == expected_hash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "check_hash(CKPT_PATH, \"e8a24ac58b8e34bdce10e0024d507f2e\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Params\n",
    "PATCH_SIZE = 16\n",
    "EMBED_DIM = 192\n",
    "RETURN_ALL_TOKENS = False\n",
    "MAX_NUMBER_CHANNELS = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load State Dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChAdaViT(\n",
    "    patch_size=PATCH_SIZE,\n",
    "    embed_dim=EMBED_DIM,\n",
    "    return_all_tokens=RETURN_ALL_TOKENS,\n",
    "    max_number_channels=MAX_NUMBER_CHANNELS,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChAdaViT(\n",
       "  (token_learner): TokenLearner(\n",
       "    (proj): Conv2d(1, 192, kernel_size=(16, 16), stride=(16, 16))\n",
       "  )\n",
       "  (pos_drop): Dropout(p=0.0, inplace=False)\n",
       "  (blocks): ModuleList(\n",
       "    (0-11): 12 x TransformerEncoderLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=192, out_features=192, bias=True)\n",
       "      )\n",
       "      (linear1): Linear(in_features=192, out_features=2048, bias=True)\n",
       "      (dropout): Dropout(p=0.0, inplace=False)\n",
       "      (linear2): Linear(in_features=2048, out_features=192, bias=True)\n",
       "      (norm1): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
       "      (norm2): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout1): Dropout(p=0.0, inplace=False)\n",
       "      (dropout2): Dropout(p=0.0, inplace=False)\n",
       "    )\n",
       "  )\n",
       "  (norm): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
       "  (head): Identity()\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert (\n",
    "    CKPT_PATH.endswith(\".ckpt\")\n",
    "    or CKPT_PATH.endswith(\".pth\")\n",
    "    or CKPT_PATH.endswith(\".pt\")\n",
    ")\n",
    "state = torch.load(CKPT_PATH, map_location=\"cpu\")[\"state_dict\"]\n",
    "for k in list(state.keys()):\n",
    "    if \"encoder\" in k:\n",
    "        state[k.replace(\"encoder\", \"backbone\")] = state[k]\n",
    "    if \"backbone\" in k:\n",
    "        state[k.replace(\"backbone.\", \"\")] = state[k]\n",
    "    del state[k]\n",
    "model.load_state_dict(state, strict=False)\n",
    "model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Random Images (Optional)\n",
    "If you are here, you probably want to test the model with your own images :)      \n",
    "But anyway, you can use the following code to generate random images with different number of channels to simply check if the model is working as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data(num_images: int, max_num_channels=MAX_NUMBER_CHANNELS):\n",
    "    imgs = []\n",
    "    labels = []\n",
    "    for i in range(num_images):\n",
    "        num_channels = np.random.randint(1, max_num_channels + 1)\n",
    "        imgs.append(torch.randn(num_channels, 224, 224))\n",
    "        labels.append(torch.randint(0, 1, (1,)))\n",
    "    data = list(zip(imgs, labels))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = generate_data(num_images=10, max_num_channels=MAX_NUMBER_CHANNELS)\n",
    "imgs, labels = zip(*data)\n",
    "distribution = {}\n",
    "for img in imgs:\n",
    "    num_channels = img.shape[0]\n",
    "    distribution[num_channels] = distribution.get(num_channels, 0) + 1\n",
    "print(\n",
    "    f\"Number of generated images: {len(imgs)} \\n Distribution of number of channels: {distribution}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the key elements of the ChAdaViT model is the ability to adapt to different number of channels. In this section, we will prepare the data to be fed into the model. We will use the `torchvision` library to load the data, and then we will create a custom dataset that will adapt the images to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_images(batch: list):\n",
    "    \"\"\"\n",
    "    Collate a batch of images into a list of channels, a list of labels and a mapping of the number of channels per image.\n",
    "\n",
    "    Args:\n",
    "        batch (list): A list of tuples of (img, label)\n",
    "\n",
    "    Return:\n",
    "        channels_list (torch.Tensor): A tensor of shape (X*num_channels, 1, height, width)\n",
    "        labels_list (torch.Tensor): A tensor of shape (batch_size, )\n",
    "        num_channels_list (list): A list of the number of channels per image\n",
    "    \"\"\"\n",
    "    num_channels_list = []\n",
    "    channels_list = []\n",
    "    labels_list = []\n",
    "\n",
    "    # Iterate over the list of images and extract the channels\n",
    "    for image, label in batch:\n",
    "        labels_list.append(label)\n",
    "        num_channels = image.shape[0]\n",
    "        num_channels_list.append(num_channels)\n",
    "\n",
    "        for channel in range(num_channels):\n",
    "            channel_image = image[channel, :, :].unsqueeze(0)\n",
    "            channels_list.append(channel_image)\n",
    "\n",
    "    channels_list = torch.cat(channels_list, dim=0).unsqueeze(\n",
    "        1\n",
    "    )  # Shape: (X*num_channels, 1, height, width)\n",
    "\n",
    "    batched_labels = torch.tensor(labels_list)\n",
    "\n",
    "    return channels_list, batched_labels, num_channels_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "collated_batch = collate_images(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "collated_batch[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def extract_features(\n",
    "    model: nn.Module,\n",
    "    batch: torch.Tensor,\n",
    "    mixed_channels: bool,\n",
    "    return_all_tokens: bool,\n",
    "):\n",
    "    \"\"\"\n",
    "    Forwards a batch of images X and extracts the features from the backbone.\n",
    "\n",
    "    Args:\n",
    "        model (nn.Module): The model to forward the images through.\n",
    "        X (torch.Tensor): The input tensor of shape (batch_size, 1, height, width).\n",
    "        list_num_channels (list): A list of the number of channels per image.\n",
    "        index (int): The index of the image to extract the features from.\n",
    "        mixed_channels (bool): Whether the images have mixed number of channels or not.\n",
    "        return_all_tokens (bool): Whether to return all tokens or not.\n",
    "\n",
    "    Returns:\n",
    "        feats (Dict): A dictionary containing the extracted features.\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "\n",
    "    # Overwrite model \"mixed_channels\" parameter for evaluation on \"normal\" datasets with uniform channels size\n",
    "    model.mixed_channels = mixed_channels\n",
    "\n",
    "    X, targets, list_num_channels = batch\n",
    "    X = X.to(device, non_blocking=True)\n",
    "\n",
    "    feats = model(x=X, index=0, list_num_channels=[list_num_channels])\n",
    "\n",
    "    if not mixed_channels:\n",
    "        if return_all_tokens:\n",
    "            # Concatenate feature embeddings per image\n",
    "            chunks = feats.view(sum(list_num_channels), -1, feats.shape[-1])\n",
    "            chunks = torch.split(chunks, list_num_channels, dim=0)\n",
    "            # Concatenate the chunks along the batch dimension\n",
    "            feats = torch.stack(chunks, dim=0)\n",
    "        # Assuming tensor is of shape (batch_size, num_tokens, backbone_output_dim)\n",
    "        feats = feats.flatten(start_dim=1)\n",
    "\n",
    "    return feats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extracted_features = extract_features(\n",
    "    model=model,\n",
    "    batch=collated_batch,\n",
    "    mixed_channels=True,\n",
    "    return_all_tokens=RETURN_ALL_TOKENS,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert extracted_features.shape[0] == len(\n",
    "    collated_batch[2]\n",
    ")  # num_embeddings == num_images, even with different number of channels\n",
    "print(\n",
    "    f\"{extracted_features.shape[0]} embeddings of dim {extracted_features.shape[1]} were extracted.\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chada-vit-dmL_hUKa-py3.10",
   "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}