Extract_SAM_encoder_weigths.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wcFP_8aRg8FT",
        "outputId": "875eac89-88ae-4485-ba47-1937eaad8b8a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "PyTorch version: 2.0.0+cu118\n",
            "Torchvision version: 0.15.1+cu118\n",
            "CUDA is available: True\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: opencv-python in /usr/local/lib/python3.9/dist-packages (4.7.0.72)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.9/dist-packages (3.7.1)\n",
            "Requirement already satisfied: numpy>=1.17.0 in /usr/local/lib/python3.9/dist-packages (from opencv-python) (1.22.4)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (0.11.0)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (4.39.3)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (23.1)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (1.0.7)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (8.4.0)\n",
            "Requirement already satisfied: importlib-resources>=3.2.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (5.12.0)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (1.4.4)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (2.8.2)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib) (3.0.9)\n",
            "Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.9/dist-packages (from importlib-resources>=3.2.0->matplotlib) (3.15.0)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.9/dist-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting git+https://github.com/facebookresearch/segment-anything.git\n",
            "  Cloning https://github.com/facebookresearch/segment-anything.git to /tmp/pip-req-build-ps4burg0\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/facebookresearch/segment-anything.git /tmp/pip-req-build-ps4burg0\n",
            "  Resolved https://github.com/facebookresearch/segment-anything.git to commit 567662b0fd33ca4b022d94d3b8de896628cd32dd\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "mkdir: cannot create directory ‘images’: File exists\n",
            "--2023-04-26 09:34:25--  https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 271475 (265K) [image/jpeg]\n",
            "Saving to: ‘images/truck.jpg.2’\n",
            "\n",
            "truck.jpg.2         100%[===================>] 265.11K  --.-KB/s    in 0.03s   \n",
            "\n",
            "2023-04-26 09:34:25 (8.46 MB/s) - ‘images/truck.jpg.2’ saved [271475/271475]\n",
            "\n",
            "--2023-04-26 09:34:25--  https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/groceries.jpg\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 168066 (164K) [image/jpeg]\n",
            "Saving to: ‘images/groceries.jpg.2’\n",
            "\n",
            "groceries.jpg.2     100%[===================>] 164.13K  --.-KB/s    in 0.02s   \n",
            "\n",
            "2023-04-26 09:34:25 (7.04 MB/s) - ‘images/groceries.jpg.2’ saved [168066/168066]\n",
            "\n",
            "--2023-04-26 09:34:25--  https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth\n",
            "Resolving dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)... 13.249.141.108, 13.249.141.13, 13.249.141.9, ...\n",
            "Connecting to dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)|13.249.141.108|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2564550879 (2.4G) [binary/octet-stream]\n",
            "Saving to: ‘sam_vit_h_4b8939.pth.2’\n",
            "\n",
            "sam_vit_h_4b8939.pt 100%[===================>]   2.39G   244MB/s    in 11s     \n",
            "\n",
            "2023-04-26 09:34:36 (233 MB/s) - ‘sam_vit_h_4b8939.pth.2’ saved [2564550879/2564550879]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torchvision\n",
        "print(\"PyTorch version:\", torch.__version__)\n",
        "print(\"Torchvision version:\", torchvision.__version__)\n",
        "print(\"CUDA is available:\", torch.cuda.is_available())\n",
        "import sys\n",
        "!{sys.executable} -m pip install opencv-python matplotlib\n",
        "!{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'\n",
        "\n",
        "!mkdir images\n",
        "!wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg\n",
        "!wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/groceries.jpg\n",
        "\n",
        "!wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5KS1bQ7Wrupl",
        "outputId": "2ffd1432-5c3a-4605-bf3a-95f7e1297d3e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: monai in /usr/local/lib/python3.9/dist-packages (1.1.0)\n",
            "Requirement already satisfied: einops in /usr/local/lib/python3.9/dist-packages (0.6.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.9/dist-packages (from monai) (1.22.4)\n",
            "Requirement already satisfied: torch>=1.8 in /usr/local/lib/python3.9/dist-packages (from monai) (2.0.0+cu118)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.9/dist-packages (from torch>=1.8->monai) (3.1.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from torch>=1.8->monai) (3.11.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch>=1.8->monai) (4.5.0)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.9/dist-packages (from torch>=1.8->monai) (3.1)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.9/dist-packages (from torch>=1.8->monai) (1.11.1)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.9/dist-packages (from torch>=1.8->monai) (2.0.0)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.9/dist-packages (from triton==2.0.0->torch>=1.8->monai) (16.0.1)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.9/dist-packages (from triton==2.0.0->torch>=1.8->monai) (3.25.2)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.9/dist-packages (from jinja2->torch>=1.8->monai) (2.1.2)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.9/dist-packages (from sympy->torch>=1.8->monai) (1.3.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install monai einops"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "atGJaFgJmSU0"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "import cv2\n",
        "import monai"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Khx05O6nDos"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "sys.path.append(\"..\")\n",
        "from segment_anything import sam_model_registry, SamPredictor\n",
        "\n",
        "sam_checkpoint = \"sam_vit_h_4b8939.pth\"\n",
        "model_type = \"vit_h\"\n",
        "\n",
        "device = \"cuda\"\n",
        "\n",
        "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
        "sam.to(device=device)\n",
        "\n",
        "predictor = SamPredictor(sam)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HIHCjWQA2_4W"
      },
      "outputs": [],
      "source": [
        "torch.save(sam.image_encoder.state_dict(),'/.models_SAM/sam_image_encoder.pth')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jg3yCswc3PPA"
      },
      "outputs": [],
      "source": [
        "from segment_anything.modeling.image_encoder import ImageEncoderViT\n",
        "from functools import partial\n",
        "from segment_anything.utils.transforms import ResizeLongestSide"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QZChItq--bUr"
      },
      "outputs": [],
      "source": [
        "prompt_embed_dim = 256\n",
        "image_size = 1024\n",
        "vit_patch_size = 16\n",
        "encoder_depth=32\n",
        "encoder_embed_dim=1280\n",
        "encoder_num_heads=16\n",
        "encoder_global_attn_indexes=[7, 15, 23, 31]\n",
        "\n",
        "image_encoder=ImageEncoderViT(\n",
        "            depth=encoder_depth,\n",
        "            embed_dim=encoder_embed_dim,\n",
        "            img_size=image_size,\n",
        "            mlp_ratio=4,\n",
        "            norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),\n",
        "            num_heads=encoder_num_heads,\n",
        "            patch_size=vit_patch_size,\n",
        "            qkv_bias=True,\n",
        "            use_rel_pos=True,\n",
        "            global_attn_indexes=encoder_global_attn_indexes,\n",
        "            window_size=14,\n",
        "            out_chans=prompt_embed_dim,\n",
        "        )\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_e6nxX0d5Bwd",
        "outputId": "9ebba4b9-68f2-4681-befe-87b772504368"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model_state=torch.load('sam_image_encoder.pth')\n",
        "image_encoder.load_state_dict(state_dict=model_state)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jMXjwuPJrmcr"
      },
      "outputs": [],
      "source": [
        "model = monai.networks.nets.UNETR(\n",
        "    in_channels=3,\n",
        "    out_channels=2,\n",
        "    img_size=1024,\n",
        "    feature_size=16,\n",
        "    hidden_size=768,\n",
        "    mlp_dim=3072,\n",
        "    num_heads=24,\n",
        "    pos_embed=\"conv\",\n",
        "    norm_name=\"batch\",\n",
        "    res_block=True,\n",
        "    dropout_rate=0.15,\n",
        "    spatial_dims=2,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aHXdUu1oCnXg",
        "outputId": "c2115541-0ee9-4e87-c9c2-7b2ae2558c59"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "ViT(\n",
              "  (patch_embedding): PatchEmbeddingBlock(\n",
              "    (patch_embeddings): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))\n",
              "    (dropout): Dropout(p=0.15, inplace=False)\n",
              "  )\n",
              "  (blocks): ModuleList(\n",
              "    (0-11): 12 x TransformerBlock(\n",
              "      (mlp): MLPBlock(\n",
              "        (linear1): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        (linear2): Linear(in_features=3072, out_features=768, bias=True)\n",
              "        (fn): GELU(approximate='none')\n",
              "        (drop1): Dropout(p=0.15, inplace=False)\n",
              "        (drop2): Dropout(p=0.15, inplace=False)\n",
              "      )\n",
              "      (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
              "      (attn): SABlock(\n",
              "        (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
              "        (qkv): Linear(in_features=768, out_features=2304, bias=False)\n",
              "        (input_rearrange): Rearrange('b h (qkv l d) -> qkv b l h d', qkv=3, l=24)\n",
              "        (out_rearrange): Rearrange('b h l d -> b l (h d)')\n",
              "        (drop_output): Dropout(p=0.15, inplace=False)\n",
              "        (drop_weights): Dropout(p=0.15, inplace=False)\n",
              "      )\n",
              "      (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
              "    )\n",
              "  )\n",
              "  (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
              ")"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model.vit"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_1yf9dJaCrD_",
        "outputId": "53d81daa-b9aa-457e-afab-e379037c5c9d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "ImageEncoderViT(\n",
              "  (patch_embed): PatchEmbed(\n",
              "    (proj): Conv2d(3, 1280, kernel_size=(16, 16), stride=(16, 16))\n",
              "  )\n",
              "  (blocks): ModuleList(\n",
              "    (0-31): 32 x Block(\n",
              "      (norm1): LayerNorm((1280,), eps=1e-06, elementwise_affine=True)\n",
              "      (attn): Attention(\n",
              "        (qkv): Linear(in_features=1280, out_features=3840, bias=True)\n",
              "        (proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
              "      )\n",
              "      (norm2): LayerNorm((1280,), eps=1e-06, elementwise_affine=True)\n",
              "      (mlp): MLPBlock(\n",
              "        (lin1): Linear(in_features=1280, out_features=5120, bias=True)\n",
              "        (lin2): Linear(in_features=5120, out_features=1280, bias=True)\n",
              "        (act): GELU(approximate='none')\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (neck): Sequential(\n",
              "    (0): Conv2d(1280, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "    (1): LayerNorm2d()\n",
              "    (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "    (3): LayerNorm2d()\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "image_encoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ngMcrHzcOUtc"
      },
      "outputs": [],
      "source": [
        "from torch.nn import functional as F\n",
        "def preprocess(x: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Normalize pixel values and pad to a square input.\"\"\"\n",
        "    # Normalize colors\n",
        "    #x = (x - self.pixel_mean) / self.pixel_std\n",
        "\n",
        "    # Pad\n",
        "    h, w = x.shape[-2:]\n",
        "    padh = image_encoder.img_size - h\n",
        "    padw = image_encoder.img_size - w\n",
        "    x = F.pad(x, (0, padw, 0, padh))\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qHj4CfXEELwY"
      },
      "outputs": [],
      "source": [
        "import torchvision.transforms as T\n",
        "image = cv2.imread('images/truck.jpg')\n",
        "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "transform=ResizeLongestSide(image_encoder.img_size)\n",
        "\n",
        "input_image = transform.apply_image(image)\n",
        "input_image_torch = torch.as_tensor(input_image,dtype=torch.float32)\n",
        "input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]\n",
        "image_transformed=preprocess(input_image_torch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oHV5fLObPAe7",
        "outputId": "418358f6-9de9-42cb-bf86-87be54d8c631"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Size([1, 3, 1024, 1024])"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "image_transformed.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ui5MAdR3ENH-"
      },
      "outputs": [],
      "source": [
        "image_encoder(image_transformed)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wUwZAJaPn6kr",
        "outputId": "09dd0c1e-84d3-4552-b847-609a5fc565ec"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Size([1, 256, 64, 64])"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "predictor.get_image_embedding().shape"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}