Day_1_SAM_Concept_Cove/Day_1_SAM_Concept_Cove.ipynb

{"cells":[{"cell_type":"markdown","metadata":{"id":"jP3-lrocLl4m"},"source":["# Navigating the Multimodal Map: Insights into Foundation Models\n","\n","## Day 1: Segment Anything Model (SAM) Concept Cove: Hands-on with multimodal models\n","\n","#### Author: Ricardo Mokhtari\n","\n","#### Course instructors: Sylwia Majchrowska and Ricardo Mokhtari, Centre for AI, DS&AI, BioPharma R&D, AstraZeneca.\n","\n","![LOGO PLACEHOLDER](https://raw.githubusercontent.com/AstraZeneca/multimodal-python-course/main/img/logo.jpg?_sm_au_=iVV6FPN3JnHkMf1ZpGsWvKttvN1NG)\n","\n","## Description of the hands on session\n","\n","In this session we will be looking at multi-modal machine learning from the point of view of the model. Specifically, we will be using a multi modal model that performs the following task:\n","\n","* *Given a input image (for example, an image of a bowl of fruits on a table) and a text prompt (for example, \"strawberries\"), the model will output the segmentation masks of only the strawberries*.\n","\n","There now exist many models capable of solving this task. However, for the purposes of this workshop, we will be working with LangSAM. Specically, we will explore 3 concepts of increasing technical challenge:\n","\n","\n","**1. Using LangSAM to segment natural images: deep dive into multimodal data fusion**\n","\n","**2. Using LangSAM for medical images**\n","\n","**3. (Exercise) Combining LangSAM and medSAM for improved medical image segmentation**\n","\n","## Navigating this notebook\n","\n","*This notebook is designed to be self-contained. Therefore, there are long descriptions in the markdown cells. You do not need to read all of the information here during the hadns on session. Instead, focus on running/modifying the code, and after the session detailed explanations are provided for you to go deeper into the topics.*\n","\n","___"]},{"cell_type":"markdown","metadata":{"id":"9nari6b0oy2Z"},"source":["# Part 1: Using LangSAM to segment natural images\n","\n","In part 1 of this notebook, we are going to learn how to use the LangSAM model to segment everyday objects in images.\n","\n","# LangSAM = Grounding DiNO + SAM\n","\n","[LangSAM](https://github.com/luca-medeiros/lang-segment-anything/tree/main) is an extention to [SAM](https://github.com/facebookresearch/segment-anything) that allows us to segment objects in image by giving text prompts. This is an extension to SAM because SAM can only segment specific objects in images when given a point or bounding box prompt.\n","\n","LangSAM achieves this by combing SAM with another model, called [Grounding DiNO](https://github.com/IDEA-Research/GroundingDINO/tree/main). Grounding DiNO takes in an image and text prompt, and predicts bounding boxes for the objects in the text prompt. These bounding boxes are then passed to SAM as input, to get the final segmentation masks. It's a neat way of adding text prompt segmentation to SAM. See the figure below.\n","\n","While both Grounding DiNO and SAM are multimodal, since they handle images + one other input modality, for the purposes of this notebook, we will be looking into the architecture of Grounding DiNO, as this is where the language/image fusion happens. The following 2 sections go into these 2 models in a bit more detail.\n","\n","![LANGSAM PLACEHOLDER](https://raw.githubusercontent.com/AstraZeneca/multimodal-python-course/main/img/LangSAM.png?_sm_au_=iVV6FPN3JnHkMf1ZpGsWvKttvN1NG)\n","\n","## Grounding DiNO\n","\n","At a high level, Grounding DiNO is a model that takes in an image and a text prompt, and returns bounding boxes corresponding to the objects mentioned in the text prompt. A detailed view at how Grounding DiNO combines text and images is shown below.\n","\n","While the below figure is quite complicated, we will zoom in on the key multimodal aspects of Grounding DiNO. Typically, for multimodal models combining images and text/other prompts, a dual-encoder + shared decoder setup is very common and has proved to work very well. You can see this below with the separate text/image backbone encoders adn the shared cross-modality decoder.\n","\n","Later in this notebook, we will dig deeper into the underlying mechanism that facilitates the fusion of modalities - cross attention.\n","\n","![GROUNDING DINO PLACEHOLDER](https://raw.githubusercontent.com/AstraZeneca/multimodal-python-course/main/img/grounding_dino.png?_sm_au_=iVV6FPN3JnHkMf1ZpGsWvKttvN1NG)\n","Image Credit: https://arxiv.org/pdf/2303.05499\n","\n","\n","## SAM\n","\n","At a high level, Segment Anything Model (SAM) is a model that takes in an image and a prompt in the form of a point or a bounding box and produces valid segmentation masks for objects defined by the prompt. An overview of the SAM model is shown below.\n","\n","Again here we see the common dual encoder, shared decoder setup that is common in multimodal models. We will not dive deeper into the SAM architecture, but cross-attention is used in a very similar way to Grounding DiNO to fuse data from the 2 modalities.\n","\n","![SAM PLACEHOLDER](https://raw.githubusercontent.com/AstraZeneca/multimodal-python-course/main/img/SAM.png?_sm_au_=iVV6FPN3JnHkMf1ZpGsWvKttvN1NG)\n","Image Credit: https://arxiv.org/abs/2304.02643\n","\n"]},{"cell_type":"markdown","metadata":{"id":"r_wZumecvQPD"},"source":["# PLEASE RUN THE CELL BELOW TO COMPLETE THE SETUP (should take around 5 minutes)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"fod5IGEKLkZC"},"outputs":[],"source":["# first, we need to clone and install LangSAM\n","!git clone https://github.com/AstraZeneca/multimodal-python-course\n","!git clone https://github.com/luca-medeiros/lang-segment-anything\n","!pip install -U git+https://github.com/luca-medeiros/lang-segment-anything.git\n","\n","# imports\n","from PIL import Image\n","from lang_sam import LangSAM\n","import matplotlib.pyplot as plt\n","import torch\n","\n","device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n","\n","# initialise the LangSAM model\n","model = LangSAM()"]},{"cell_type":"markdown","metadata":{"id":"Gi-egyrpsJzc"},"source":["## Below is the code for initialising the main langSAM functions"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"bHdhOSFJU-FM"},"outputs":[],"source":["def display_image_with_masks(image, masks):\n","    \"\"\"Function for plotting the predicted masks from SAM\n","\n","    image (PIL.Image): input image as PIL.Image RGB\n","    masks: output masks from LangSAM\n","    \"\"\"\n","    # get the number of masks\n","    num_masks = len(masks)\n","    # if there were no predicted masks, exit\n","    if num_masks == 0:\n","        print(\"no masks found\")\n","        return\n","\n","    # create figure and plot\n","    fig, axes = plt.subplots(1, num_masks + 1, figsize=(15, 5))\n","    axes[0].imshow(image)\n","    axes[0].set_title(\"Original Image\")\n","    axes[0].axis('off')\n","\n","    # plot each mask as a new figure\n","    # TODO:\n","    for i, mask_np in enumerate(masks):\n","        axes[i+1].imshow(mask_np, cmap='gray')\n","        axes[i+1].set_title(f\"Mask {i+1}\")\n","        axes[i+1].axis('off')\n","\n","    plt.tight_layout()\n","    plt.show()\n","\n","def display_image_with_boxes(image, boxes, logits):\n","    \"\"\"Function for overlaying predicted bounding boxes on the image\n","\n","    image (PIL.Image): input image as PIL.Image RGB\n","    boxes (np.array): coordinates of the bounding boxes\n","    logits (torch.tensor): the logits for each bounding box prediction\n","    \"\"\"\n","    # creat the figure\n","    fig, ax = plt.subplots(dpi=150)\n","    ax.imshow(image)\n","    ax.set_title(\"Image with Bounding Boxes\")\n","    ax.axis('off')\n","\n","    # for each box and logit\n","    for box, logit in zip(boxes, logits):\n","        # getting the coordinates\n","        x_min, y_min, x_max, y_max = box\n","        confidence_score = round(logit.item(), 2)  # Convert logit to a scalar before rounding\n","        box_width = x_max - x_min\n","        box_height = y_max - y_min\n","\n","        # Draw bounding box\n","        rect = plt.Rectangle((x_min, y_min), box_width, box_height, fill=False, edgecolor='red', linewidth=2)\n","        ax.add_patch(rect)\n","\n","        # Add confidence score as text\n","        ax.text(x_min, y_min, f\"Confidence: {confidence_score}\", fontsize=8, color='red', verticalalignment='top')\n","\n","    plt.show()"]},{"cell_type":"markdown","metadata":{"id":"-eG-yIbutaKl"},"source":["### In the cell below, you can play around with LangSAM. You can replace the image inputted by giving a new path to image_pil. You can edit the text prompt to segment different objects.\n","\n","### You can adjust some of the LangSAM parameters to fine-tune your prediction. \"box_threshold\" defines the confidence needed for the bounding box prediction to be displayed - lowering this value means you will get more prediction outuputs\n","\n","### There are 5 different images in the inference_images/natural_images folder. These images are from the [COCO dataset](https://cocodataset.org/#home)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"ct9VU8C5s8OO"},"outputs":[],"source":["# input your image here - there are 5 images each with filepath X.jpg\n","file_path = \"multimodal-python-course/img/inference_images/natural_images/1.jpg\"\n","image_pil = Image.open(file_path).convert(\"RGB\")\n","\n","# input your text prompt here - tip: if you want to segment multiple different\n","# objects, try separating them with a full stop.\n","text_prompt = \"cat\"\n","\n","# run inference\n","masks, boxes, phrases, logits = model.predict(image_pil, text_prompt, box_threshold=0.35)\n","\n","# display the predictions\n","display_image_with_boxes(image_pil, boxes, logits)\n","display_image_with_masks(image_pil, masks)"]},{"cell_type":"markdown","metadata":{"id":"3WSneYelxe2_"},"source":["## In the next cell we will look at how the text-image fusion is performed. In Groudning DiNO, this is done using cross attention.\n","\n","#### What is attention?\n","\n","![SELF ATTENTION PLACEHOLDER](https://raw.githubusercontent.com/AstraZeneca/multimodal-python-course/main/img/self_attention.png?_sm_au_=iVV6FPN3JnHkMf1ZpGsWvKttvN1NG)\n","Image credit: Raimi Karim, https://towardsdatascience.com/illustrated-self-attention-2d627e33b20a\n","\n","\"Attention\" generally in machine learning is a concept that originated in natural language processing (NLP). In language, a single word can have multiple meanings depending on the context. For example, the word \"great\" is generally positive. However, depending on the context, for example \"that trip was a great waste of my time\" it can have a negative context. The context length can also vary greatly. For example, \"the trees were so tall they stretched hundreds of metres in the air and were so wide that you could not wrap your arms around them.\". In this sentence, even though there are many nouns, the word \"them\" refers specically to the word \"trees\" and no other word. Yet the word \"trees\" is located very far away in the sentence from the word \"them\".\n","\n","Attention is a modelling approach to capture these inter-word relations. The same concepts are true of images - the spatial organisation of objects can span the entire scene and an object's context can change the meaning of the scene. More specifically, an attention mechanism will process a sequence of inputs (for example words or image patches) and capture how strongly each word is related to all other words. In this way, attention models words in their context over the entire sequence (sentence of whole image).\n","\n","The most common form of attention is self-attention, where the importance of a given input (eg. word) to all other inputs, including itself, are calculated. For a deep dive into how self-attention works, please look at this [excellent blog](https://towardsdatascience.com/illustrated-self-attention-2d627e33b20a).\n","\n","#### What is cross attention?\n","\n","![CROSS ATTENTION PLACEHOLDER](https://raw.githubusercontent.com/AstraZeneca/multimodal-python-course/main/img/cross_attention.png?_sm_au_=iVV6FPN3JnHkMf1ZpGsWvKttvN1NG)\n","Image credit: Vaclav Kosar, https://vaclavkosar.com/ml/cross-attention-in-transformer-architecture\n","\n","\n","Cross attention is a specific type of the general attention mechanism. It is slightly different from the more commonly known self-attention that powers models like ChatGPT and other transformers. Self-attention works based on a single input sequence. The sequence (where each token is transformed into an embedding) is fed into a self attention module. The self attention module calculates how important each entry in the sequence is to all other entries, including itself.\n","\n","However, in cross attention, we don't just have one sequence, we have 2. For example, we have a sequence of text features, and a sequence of image features. Cross attention allows to compute the importance of all text features to all image features, and the other way round. In the end, we transform the text and image features to new features that capture information from both modalities.\n","\n","#### Understanding the model's embeddings\n","\n","Plotting the emebddings of a model can be a very useful tool for understanding what the model has learnt. By plotting the embeddings, we can start to explore the models latent space and understand what features of different data points the model considers to be similar. For example, if we plot embeddings for all our data points and we see that cats and dogs are grouped together in one cluster and that cars and trucks are grouped into another cluster, then we can have some confidence that the model has learnt that animals and vehicles have distinct features and that the model is able to distinguish them.\n","\n","On Day 2, you will see a real world example of embeddings plotting and how it can be useful. Below we have provided a toy example for completeness, although the embeddings below are only randomly generated.\n","\n","For a great and short video looking into these concepts in more detail, click [here](https://www.youtube.com/watch?v=aw3H-wPuRcw&list=WL&index=36)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"cKISKMzimahW"},"outputs":[],"source":["# cross attention toy example\n","import torch\n","import torch.nn as nn\n","\n","# change here if you want to see slightly different outputs\n","torch.manual_seed(0)\n","\n","def plot_embeddings(x1, x2, output):\n","    \"\"\"Functio for plotting embeddings\n","    x1 - the embeddings of sequence 1\n","    x2 - the embeddings of sequence 2\n","    output - the emebeddings after cross attention\n","    \"\"\"\n","    # we can only plot 2 dimensional embeddings\n","    if x1.shape[2] == 2 and x2.shape[2] == 2 and output.shape[2] == 2:\n","\n","        # get the x and y of all embeddings\n","        x1_x = x1.numpy()[0,:,0]\n","        x1_y = x1.numpy()[0,:,1]\n","\n","        x2_x = x2.numpy()[0,:,0]\n","        x2_y = x2.numpy()[0,:,1]\n","\n","        output_x = output.numpy()[0,:,0]\n","        output_y = output.numpy()[0,:,1]\n","\n","        # create plot\n","        plt.figure()\n","        plt.scatter(x1_x, x1_y, label=\"x1\")\n","        plt.scatter(x2_x, x2_y, label=\"x2\")\n","        plt.scatter(output_x, output_y, label=\"output\")\n","        plt.xlabel(\"Emebedding dim 1\")\n","        plt.ylabel(\"Embedding dim 2\")\n","        plt.legend()\n","        plt.tight_layout()\n","        plt.show()\n","    else:\n","        print(\"The dimensions of the embeddings is not 2, these can't be plottted\")\n","        return\n","\n","class CrossAttention(nn.Module):\n","    def __init__(self, input_dim, hidden_dim):\n","        \"\"\"Initialistion toy cross attention example\n","        input_dim - the dimensions of the input embeddings\n","        hidden_dim - the dimension of the attention blocks\n","        \"\"\"\n","        super(CrossAttention, self).__init__()\n","\n","        self.input_dim = input_dim\n","        self.hidden_dim = hidden_dim\n","\n","        # Linear transformation for query, key, and value\n","        self.linear_query = nn.Linear(input_dim, hidden_dim)\n","        self.linear_key = nn.Linear(input_dim, hidden_dim)\n","        self.linear_value = nn.Linear(input_dim, hidden_dim)\n","\n","        # Softmax layer for attention weights\n","        self.softmax = nn.Softmax(dim=-1)\n","\n","    def forward(self, x1, x2):\n","        # x1: (batch_size, seq_len1, input_dim)\n","        # x2: (batch_size, seq_len2, input_dim)\n","\n","        print(\"First, we create the inputs, which are 2 sequences of embeddings\")\n","        print(f\"Sequence 1 (x1) has shape [batch size (BS), sequence 1 length (SL1), input size (IS)]: {x1.shape}\")\n","        print(f\"x1:\\n {x1}\")\n","        print(f\"Sequence 2 (x2) has shape [batch size (BS), sequence 2 length (SL2), input size (IS)]: {x2.shape}\")\n","        print(f\"x2:\\n {x2}\\n\")\n","\n","        # Project input vectors to the hidden space\n","        print(\"First we transform x1 to the query space\")\n","        query = self.linear_query(x1)  # (batch_size, seq_len1, hidden_dim)\n","        print(f\"Shape of query (Q) [BS, SL1, hidden dim (HD)]: {query.shape}\\n\")\n","\n","        print(\"Next we transform x2 to the key space\")\n","        key = self.linear_key(x2)      # (batch_size, seq_len2, hidden_dim)\n","        print(f\"Shape of key (K) [BS, SL2, hidden dim (HD)]: {key.shape}\\n\")\n","\n","        print(\"Next we transform x2 to the value space\")\n","        value = self.linear_value(x2)  # (batch_size, seq_len2, hidden_dim)\n","        print(f\"Shape of value (V) [BS, SL2, hidden dim (HD)]: {value.shape}\\n\")\n","\n","        # Compute attention scores\n","        print(\"Next, we compute the dot product attention between the query and key (transposed)\")\n","        attn_scores = torch.bmm(query, key.transpose(1, 2))  # (batch_size, seq_len1, seq_len2)\n","        print(f\"Shape of query: {query.shape}, Shape of transposed key: {key.transpose(1,2).shape}\")\n","        print(f\"Shape of attention matrix [BS, SL1, SL2]: {attn_scores.shape}\")\n","        print(f\"Raw attention matrix:\\n {attn_scores}\\n\")\n","\n","        # Apply softmax to get attention weights\n","        print(\"Next, we simply apply the softmax to the attention matrix\")\n","        attn_weights = self.softmax(attn_scores)  # (batch_size, seq_len1, seq_len2)\n","        print(f\"Softmaxed attention matrix:\\n {attn_weights}\\n\")\n","\n","        # Apply attention weights to values\n","        print(\"Finally, we multiply the softmaxed attention weights with the value matrix\")\n","        attended_values = torch.bmm(attn_weights, value)  # (batch_size, seq_len1, hidden_dim)\n","        print(f\"Shape of output [BS, SL1, HD]: {attended_values.shape}\")\n","        print(f\"Final output:\\n {attended_values}\")\n","\n","        return attended_values\n","\n","# Set dimensions of everything\n","input_dim = 2   # the dimension of the input embeddings\n","hidden_dim = 2  # the dimension of the attention blocks\n","batch_size = 1  # the batch size (ignore)\n","seq_len1 = 3    # the length of the first sequence (number of embeddings)\n","seq_len2 = 3    # the length of the second sequence (number of embeddings)\n","\n","# remove gradients\n","with torch.no_grad():\n","    # create some dummy embeddings\n","    x1 = torch.randn(batch_size, seq_len1, input_dim)\n","    x2 = torch.randn(batch_size, seq_len2, input_dim)\n","\n","    # apply cross attention\n","    cross_attention = CrossAttention(input_dim, hidden_dim)\n","    output = cross_attention(x1, x2)\n","\n","    plot_embeddings(x1, x2, output)"]},{"cell_type":"markdown","metadata":{"id":"qLq9ZXktRx-D"},"source":["## The implementation of Cross Attention in Grounding DiNO\n","\n","Above we looked at a toy example of cross attention. Below we can see how it is implemented in the Grounding DiNO model. We can do this by extracting the weights of the part of the model that performs this operation and passing some dummy inputs through it."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Ci91N1i4C51i"},"outputs":[],"source":["import torch\n","import numpy as np\n","from groundingdino.util.misc import NestedTensor\n","import groundingdino.datasets.transforms as T\n","from groundingdino.util import get_tokenlizer\n","from groundingdino.models.GroundingDINO.backbone.swin_transformer import SwinTransformer\n","\n","def preprocess_caption(caption: str) -> str:\n","    result = caption.lower().strip()\n","    if result.endswith(\".\"):\n","        return result\n","    return result + \".\"\n","\n","# remove grounding dino model from langsam\n","grounding_dino = model.groundingdino\n","\n","with torch.no_grad():\n","\n","    # extract the first \"fusion layer\" which performs cross attention\n","    w = grounding_dino.transformer.encoder.fusion_layers[0]\n","    print(w)\n","\n","    # creating some inputs of shape [batch, sequence length, embedding size]\n","    v = torch.rand(1,28,256).to(device)\n","    print(f\"input sequence 1 v: {v.shape}\")\n","    l = torch.rand(1,28,256).to(device)\n","    print(f\"input sequence 2 l: {v.shape}\")\n","    out = w.forward(v,l)\n","    print(f\"output of fusion layer: {out[0].shape}\")"]},{"cell_type":"markdown","metadata":{"id":"OxpT7mSfoOCw"},"source":["If you are curious about playing around specifically with the image and text embeddings that come out of grounding dino, for example, for plotting purposes, you can extract the image and text embedders from grounding dino and run them separately on each input, as in the cell below."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"YDU1VwmLoM5g"},"outputs":[],"source":["# encoder texts\n","dummy_caption = preprocess_caption(\"here. is. some. text.\")\n","dummy_image = torch.rand(1,3,224,224).to(device)\n","\n","with torch.no_grad():\n","    # image encoder\n","    trans = grounding_dino.backbone[0]\n","    trans = trans.to(device)\n","    img_enc = trans.forward_raw(dummy_image.to(device))[0]\n","    print(f\"shape of image enc: {img_enc.shape}\")\n","\n","    # text encoder\n","    bert = grounding_dino.bert\n","    tokenized = grounding_dino.tokenizer(dummy_caption, padding=\"longest\", return_tensors=\"pt\").to(device)\n","    text_enc = bert(**tokenized)  # bs, 195, 768\n","    print(f\"shape of text enc: {text_enc.pooler_output.shape}\")"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"U0zVvmR8PBA0"},"outputs":[],"source":[]},{"cell_type":"markdown","metadata":{"id":"nBPRWUVbo52p"},"source":["# Part 2: Using LangSAM to segment medical images\n","\n","In this part, we are going to attempt to use LangSAM to segment objects in medical images. Since the components of LangSAM (grounding DiNO and SAM) were both trained only with natural, non medical images, it's likely that it will struggle.\n","\n","In the inference_images/medical_images folder, there are 3 images of breast ultrasounds, taken from [this dataset](https://www.kaggle.com/datasets/aryashah2k/breast-ultrasound-images-dataset). Each image has an associated mask so you can check how well each image has been segmented.\n","\n","As part of this exercise, try to probe where the problem lies with langSAM - is it with Grounding DiNO (bounding boxes produced do not reflect the medical object) or with SAM (valid segmentation masks are not produced in the bounding box region)?\n","\n","In addition, especially for the breast cancer set, try to probe how the text prompts are interpreted by Grounding DiNO. For example, does Grounding DiNO produce the same bounding box when you use different prompts other than \"tumor\"? This points towards the robustness of how well different concepts are represented in Grounding DiNO's embedding space.\n","\n","Try to see if there is any way to use prompt engineering (modifying your text prompt) to improve the outputs of either model."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"N5rhmAqRPBC0"},"outputs":[],"source":["def display_medical_image(filepath, text_prompt=\"\"):\n","    \"\"\"\n","    Function for displaying prediction and GT for medical images\n","\n","    filepath (str): path to file\n","    text_promot (str): the text prompt for Grounding dino\n","    filetype (str): \"breast_ultrasound\" or \"pathology\"\n","    \"\"\"\n","\n","    # input your image here\n","    image_pil = Image.open(filepath).convert(\"RGB\")\n","\n","    # run inference\n","    masks, boxes, phrases, logits = model.predict(image_pil, text_prompt, box_threshold=0.35)\n","\n","    # display the predictions\n","    display_image_with_boxes(image_pil, boxes, logits)\n","    display_image_with_masks(image_pil, masks)\n","\n","    try:\n","        mask_path = filepath[:-4]+\"_mask.png\"\n","\n","        fig, axs = plt.subplots(nrows=1, ncols=2)\n","        axs[0].imshow(image_pil)\n","        axs[0].set_title(\"Original image\")\n","        axs[0].axis('off')\n","        axs[1].imshow(Image.open(mask_path))\n","        axs[1].set_title(\"Ground truth\")\n","        axs[1].axis('off')\n","\n","        plt.show()\n","    except:\n","        print(\"no ground truth found!\")\n","\n","# load in a medical image - stored in inference_images/medical_images\n","display_medical_image(\n","    \"multimodal-python-course/img/inference_images/medical_images/malignant_3.png\",\n","    text_prompt=\"tumor\",\n",")"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"ueSkxKagPBFK"},"outputs":[],"source":[]},{"cell_type":"markdown","metadata":{"id":"WIIzFgBcQRQp"},"source":["# Part 3 (Exercise): Replacing SAM in LangSAM with MedSAM for improved masks\n","\n","*This is an optional exercise for those interested in continuing extending their learning on this topic.*\n","\n","In Part 2 above, we saw that LangSAM performs poorly for medical image segmentation. The question is can we do anything to improve it? The best way would be to replace either or both of Grounding DiNO and SAM with versions that were trained on medical images. For this exercise, we will focus on replacing the SAM model with MedSAM, which is a version of SAM that was trained on a large dataset of medical images.\n","\n","## Guidance\n","\n","Below we have provided some skeleton code for loading in the MedSAM model and running segmentation on an example CT image. You can clearly see from this example that MedSAM is superior to SAM for medical images, as it is able to precisely segment the rather faint object in the bounding box. In the following, we will point to the key functions and I/O that needs to be edited to replace SAM with medSAM.\n","\n","The main difficulty you will run into here is the difference between how SAM and MedSAM handle inputs and outputs - make sure that you print the shapes of all input/output variables to investigate what they represent.\n","\n","To get started - the main driving code for LangSAM is in lang-segment-anything/lang_sam.py. Below is a modified snippet of the 3 prediction functions that you will need to edit. You will see that there is a **predict_sam()** method - try to add a new method **predict_medsam()** that will load the medsam model, receive the boxes from **predict_dino()** and create the segmentation masks.\n","\n","Good luck! If you are able to complete this exercise, consider sharing your findings and learnings with the rest of the group. Do you think that replacing SAM with medSAM is enough to make LangSAM usable for medical images? Or is it more important to get good bounding boxes first/understand medical language in the first step?\n","\n","```python\n","\"\"\" this code is a snippet from from lang-segment-anything/lang_sam.py\"\"\"\n","# function for predicting the bounding boxes with grounding DiNO\n","# input is the image, text prompt and thresholds\n","def predict_dino(self, image_pil, text_prompt, box_threshold, text_threshold):\n","    image_trans = transform_image(image_pil)\n","    boxes, logits, phrases = predict(model=self.groundingdino,\n","                                        image=image_trans,\n","                                        caption=text_prompt,\n","                                        box_threshold=box_threshold,\n","                                        text_threshold=text_threshold,\n","                                        device=self.device)\n","    W, H = image_pil.size\n","    boxes = box_ops.box_cxcywh_to_xyxy(boxes) * torch.Tensor([W, H, W, H])\n","\n","    return boxes, logits, phrases\n","\n","# function for predicting with sam\n","# input is the image and the boxes that come out of the above function\n","def predict_sam(self, image_pil, boxes):\n","    image_array = np.asarray(image_pil)\n","    self.sam.set_image(image_array)\n","    transformed_boxes = self.sam.transform.apply_boxes_torch(boxes, image_array.shape[:2])\n","    masks, _, _ = self.sam.predict_torch(\n","        point_coords=None,\n","        point_labels=None,\n","        boxes=transformed_boxes.to(self.sam.device),\n","        multimask_output=False,\n","    )\n","    return masks.cpu()\n","\n","def predict_medsam(self, ...):\n","    ### Your code here ###\n","\n","# the overall prediction functiont that combines the above 2 methods\n","# input\n","def predict(self, image_pil, text_prompt, box_threshold=0.3, text_threshold=0.25, sam_model=\"SAM\"):\n","    boxes, logits, phrases = self.predict_dino(image_pil, text_prompt, box_threshold, text_threshold)\n","    masks = torch.tensor([])\n","    if len(boxes) > 0:\n","        if sam_model == \"SAM\":\n","            masks = self.predict_sam(image_pil, boxes)\n","            masks = masks.squeeze(1)\n","        else:\n","            ### Your code here ###\n","    return masks, boxes, phrases, logits\n","```"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"u_hxUDctQSSD"},"outputs":[],"source":["!pip install git+https://github.com/bowang-lab/MedSAM.git"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"bGIoezxFQYxg"},"outputs":[],"source":["# all code taken from https://github.com/bowang-lab/MedSAM\n","# %% environment and functions\n","import numpy as np\n","import matplotlib.pyplot as plt\n","import os\n","join = os.path.join\n","import torch\n","from segment_anything import sam_model_registry\n","from skimage import io, transform\n","import torch.nn.functional as F\n","\n","# visualization functions\n","# source: https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb\n","# change color to avoid red and green\n","def show_mask(mask, ax, random_color=False):\n","    if random_color:\n","        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n","    else:\n","        color = np.array([251/255, 252/255, 30/255, 0.6])\n","    h, w = mask.shape[-2:]\n","    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n","    ax.imshow(mask_image)\n","\n","def show_box(box, ax):\n","    x0, y0 = box[0], box[1]\n","    w, h = box[2] - box[0], box[3] - box[1]\n","    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='blue', facecolor=(0,0,0,0), lw=2))\n","\n","@torch.no_grad()\n","def medsam_inference(medsam_model, img_embed, box_1024, H, W):\n","    box_torch = torch.as_tensor(box_1024, dtype=torch.float, device=img_embed.device)\n","    if len(box_torch.shape) == 2:\n","        box_torch = box_torch[:, None, :] # (B, 1, 4)\n","\n","    sparse_embeddings, dense_embeddings = medsam_model.prompt_encoder(\n","        points=None,\n","        boxes=box_torch,\n","        masks=None,\n","    )\n","    low_res_logits, _ = medsam_model.mask_decoder(\n","        image_embeddings=img_embed, # (B, 256, 64, 64)\n","        image_pe=medsam_model.prompt_encoder.get_dense_pe(), # (1, 256, 64, 64)\n","        sparse_prompt_embeddings=sparse_embeddings, # (B, 2, 256)\n","        dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64)\n","        multimask_output=False,\n","        )\n","\n","    low_res_pred = torch.sigmoid(low_res_logits)  # (1, 1, 256, 256)\n","\n","    low_res_pred = F.interpolate(\n","        low_res_pred,\n","        size=(H, W),\n","        mode=\"bilinear\",\n","        align_corners=False,\n","    )  # (1, 1, gt.shape)\n","    low_res_pred = low_res_pred.squeeze().cpu().numpy()  # (256, 256)\n","    medsam_seg = (low_res_pred > 0.5).astype(np.uint8)\n","    return medsam_seg\n","\n","# download model and data\n","img_id = '1Qf0IQF1zTwOJ4Fh4Ix5rtVhR9cnNbI2M'\n","!gdown $img_id\n","model_id = \"1UAmWL88roYR7wKlnApw5Bcuzf2iQgk6_&confirm=t\"\n","!gdown $model_id\n","\n","#%% load model and image\n","MedSAM_CKPT_PATH = \"medsam_vit_b.pth\"\n","device = \"cuda:0\"\n","medsam_model = sam_model_registry['vit_b'](checkpoint=MedSAM_CKPT_PATH)\n","medsam_model = medsam_model.to(device)\n","medsam_model.eval()\n","\n","img_np = io.imread('img_demo.png')\n","if len(img_np.shape) == 2:\n","    img_3c = np.repeat(img_np[:, :, None], 3, axis=-1)\n","else:\n","    img_3c = img_np\n","H, W, _ = img_3c.shape\n","\n","#%% image preprocessing and model inference\n","img_1024 = transform.resize(img_3c, (1024, 1024), order=3, preserve_range=True, anti_aliasing=True).astype(np.uint8)\n","img_1024 = (img_1024 - img_1024.min()) / np.clip(\n","    img_1024.max() - img_1024.min(), a_min=1e-8, a_max=None\n",")  # normalize to [0, 1], (H, W, 3)\n","# convert the shape to (3, H, W)\n","img_1024_tensor = torch.tensor(img_1024).float().permute(2, 0, 1).unsqueeze(0).to(device)\n","\n","box_np = np.array([[95,255, 190, 350]])\n","# transfer box_np t0 1024x1024 scale\n","box_1024 = box_np / np.array([W, H, W, H]) * 1024\n","with torch.no_grad():\n","    image_embedding = medsam_model.image_encoder(img_1024_tensor) # (1, 256, 64, 64)\n","\n","medsam_seg = medsam_inference(medsam_model, image_embedding, box_1024, H, W)\n","\n","#%% visualize results\n","fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n","ax[0].imshow(img_3c)\n","show_box(box_np[0], ax[0])\n","ax[0].set_title(\"Input Image and Bounding Box\")\n","ax[1].imshow(img_3c)\n","show_mask(medsam_seg, ax[1])\n","show_box(box_np[0], ax[1])\n","ax[1].set_title(\"MedSAM Segmentation\")\n","plt.show()"]},{"cell_type":"markdown","metadata":{"id":"X9kuFnGrpK4r"},"source":["# That's the end of the notebook!\n","\n","Here is the list of all resources shared in the notebook and a couple extras:\n","\n","1. LangSAM [Code](https://github.com/luca-medeiros/lang-segment-anything)\n","2. Grounding DiNO [Code](https://github.com/IDEA-Research/GroundingDINO) [Paper](https://arxiv.org/abs/2303.05499)\n","3. SAM [Code](https://github.com/facebookresearch/segment-anything) [Paper](https://arxiv.org/abs/2304.02643)\n","4. [Attention illustrated blog](https://towardsdatascience.com/illustrated-self-attention-2d627e33b20a)\n","5. [Attention video](https://www.youtube.com/watch?v=KmAISyVvE1Y)\n","6. [Another attention video](https://www.youtube.com/watch?v=eMlx5fFNoYc)\n","7. [Cross attention](https://www.youtube.com/watch?v=aw3H-wPuRcw&list=WL&index=37)\n","8. [Visualise a transformer](https://bbycroft.net/llm)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"3sbae0W-T8lg"},"outputs":[],"source":[]}],"metadata":{"accelerator":"GPU","colab":{"gpuType":"T4","provenance":[]},"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.11.7"}},"nbformat":4,"nbformat_minor":0}