{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xhhr4iSQuQq_"
      },
      "source": [
        "#Efficient SAM Example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AIrAUKnLClPD"
      },
      "source": [
        "This script provides example for how to get visualization result from EfficientSAM using weight file, part of the code is borrow from MobileSAM project, many thanks!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zylNfpYIuXeR"
      },
      "source": [
        "##Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "I64YhiKsS2KU"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch\n",
        "from torchvision.transforms import ToTensor\n",
        "from PIL import Image\n",
        "import io"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pw_4lyT8uMvy"
      },
      "source": [
        "#Box and Point prompt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "4hrhmpHroFUH"
      },
      "outputs": [],
      "source": [
        "def run_ours_box_or_points(img_path, pts_sampled, pts_labels, model):\n",
        "    image_np = np.array(Image.open(img_path))\n",
        "    img_tensor = ToTensor()(image_np)\n",
        "    pts_sampled = torch.reshape(torch.tensor(pts_sampled), [1, 1, -1, 2])\n",
        "    pts_labels = torch.reshape(torch.tensor(pts_labels), [1, 1, -1])\n",
        "    predicted_logits, predicted_iou = model(\n",
        "        img_tensor[None, ...],\n",
        "        pts_sampled,\n",
        "        pts_labels,\n",
        "    )\n",
        "\n",
        "    sorted_ids = torch.argsort(predicted_iou, dim=-1, descending=True)\n",
        "    predicted_iou = torch.take_along_dim(predicted_iou, sorted_ids, dim=2)\n",
        "    predicted_logits = torch.take_along_dim(\n",
        "        predicted_logits, sorted_ids[..., None, None], dim=2\n",
        "    )\n",
        "\n",
        "    return torch.ge(predicted_logits[0, 0, 0, :, :], 0).cpu().detach().numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-83WUeNPuJnT"
      },
      "source": [
        "#Visualization Related"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "QKWt76-AG31h"
      },
      "outputs": [],
      "source": [
        "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([30 / 255, 144 / 255, 255 / 255, 0.8])\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",
        "\n",
        "def show_points(coords, labels, ax, marker_size=375):\n",
        "    pos_points = coords[labels == 1]\n",
        "    neg_points = coords[labels == 0]\n",
        "    ax.scatter(\n",
        "        pos_points[:, 0],\n",
        "        pos_points[:, 1],\n",
        "        color=\"green\",\n",
        "        marker=\"*\",\n",
        "        s=marker_size,\n",
        "        edgecolor=\"white\",\n",
        "        linewidth=1.25,\n",
        "    )\n",
        "    ax.scatter(\n",
        "        neg_points[:, 0],\n",
        "        neg_points[:, 1],\n",
        "        color=\"red\",\n",
        "        marker=\"*\",\n",
        "        s=marker_size,\n",
        "        edgecolor=\"white\",\n",
        "        linewidth=1.25,\n",
        "    )\n",
        "\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(\n",
        "        plt.Rectangle((x0, y0), w, h, edgecolor=\"yellow\", facecolor=(0, 0, 0, 0), lw=5)\n",
        "    )\n",
        "\n",
        "def show_anns_ours(mask, ax):\n",
        "    ax.set_autoscale_on(False)\n",
        "    img = np.ones((mask.shape[0], mask.shape[1], 4))\n",
        "    img[:, :, 3] = 0\n",
        "    color_mask = [0, 1, 0, 0.7]\n",
        "    img[np.logical_not(mask)] = color_mask\n",
        "    ax.imshow(img)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GHj10cGetlGN"
      },
      "source": [
        "#Create the model and load the weights from the checkpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vnzGws0DBJsj",
        "outputId": "18518a50-4b0b-4568-c6e4-e919fd7c6b44"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'EfficientSAM'...\n",
            "remote: Enumerating objects: 339, done.\u001b[K\n",
            "remote: Counting objects: 100% (55/55), done.\u001b[K\n",
            "remote: Compressing objects: 100% (18/18), done.\u001b[K\n",
            "remote: Total 339 (delta 52), reused 37 (delta 37), pack-reused 284\u001b[K\n",
            "Receiving objects: 100% (339/339), 222.75 MiB | 13.03 MiB/s, done.\n",
            "Resolving deltas: 100% (194/194), done.\n",
            "Updating files: 100% (35/35), done.\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/yformer/EfficientSAM.git\n",
        "import os\n",
        "os.chdir(\"EfficientSAM\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3HR4CCpYUpAI",
        "outputId": "f555f0c8-b5d3-49e5-af62-6babb15a00ae"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "SqueezeSam(\n",
              "  (image_encoder): UnetEncoder(\n",
              "    (proj): Sequential(\n",
              "      (0): Conv2d(5, 64, kernel_size=(2, 2), stride=(2, 2))\n",
              "      (1): GroupNorm(1, 64, eps=1e-05, affine=True)\n",
              "    )\n",
              "    (down_layers): ModuleList(\n",
              "      (0): DoubleConv(\n",
              "        (double_conv): Sequential(\n",
              "          (0): Conv2dWithActivation(\n",
              "            (conv): Sequential(\n",
              "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "              (1): ReLU()\n",
              "            )\n",
              "            (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "            (quant_output_normalization): QuantStub()\n",
              "          )\n",
              "          (1): Conv2dWithActivation(\n",
              "            (conv): Sequential(\n",
              "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "              (1): ReLU()\n",
              "            )\n",
              "            (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "            (quant_output_normalization): QuantStub()\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (1): Down(\n",
              "        (maxpool_conv): Sequential(\n",
              "          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "          (1): DoubleConv(\n",
              "            (double_conv): Sequential(\n",
              "              (0): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "              (1): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (2-4): 3 x Down(\n",
              "        (maxpool_conv): Sequential(\n",
              "          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "          (1): DoubleConv(\n",
              "            (double_conv): Sequential(\n",
              "              (0): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "              (1): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (5): Down(\n",
              "        (maxpool_conv): Sequential(\n",
              "          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "          (1): DoubleConv(\n",
              "            (double_conv): Sequential(\n",
              "              (0): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(128, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "              (1): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(164, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (6): Down(\n",
              "        (maxpool_conv): Sequential(\n",
              "          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "          (1): DoubleConv(\n",
              "            (double_conv): Sequential(\n",
              "              (0): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(164, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "              (1): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (7): Down(\n",
              "        (maxpool_conv): Sequential(\n",
              "          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "          (1): DoubleConv(\n",
              "            (double_conv): Sequential(\n",
              "              (0): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "              (1): Conv2dWithActivation(\n",
              "                (conv): Sequential(\n",
              "                  (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                  (1): ReLU()\n",
              "                )\n",
              "                (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "                (quant_output_normalization): QuantStub()\n",
              "              )\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (prompt_encoder): NullPromptEncoder()\n",
              "  (mask_decoder): MaskDecoder(\n",
              "    (transformer): TwoWayTransformer(\n",
              "      (layers): ModuleList(\n",
              "        (0-1): 2 x TwoWayAttentionBlock(\n",
              "          (self_attn): AttentionForTwoWayAttentionBlock(\n",
              "            (q_proj): Linear(in_features=256, out_features=256, bias=True)\n",
              "            (k_proj): Linear(in_features=256, out_features=256, bias=True)\n",
              "            (v_proj): Linear(in_features=256, out_features=256, bias=True)\n",
              "            (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
              "          )\n",
              "          (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "          (cross_attn_token_to_image): AttentionForTwoWayAttentionBlock(\n",
              "            (q_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "            (k_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "            (v_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "            (out_proj): Linear(in_features=128, out_features=256, bias=True)\n",
              "          )\n",
              "          (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "          (mlp): MLPBlock(\n",
              "            (layers): ModuleList(\n",
              "              (0): Sequential(\n",
              "                (0): Linear(in_features=256, out_features=256, bias=True)\n",
              "                (1): ReLU()\n",
              "              )\n",
              "            )\n",
              "            (fc): Linear(in_features=256, out_features=256, bias=True)\n",
              "          )\n",
              "          (norm3): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "          (norm4): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "          (cross_attn_image_to_token): AttentionForTwoWayAttentionBlock(\n",
              "            (q_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "            (k_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "            (v_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "            (out_proj): Linear(in_features=128, out_features=256, bias=True)\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (final_attn_token_to_image): AttentionForTwoWayAttentionBlock(\n",
              "        (q_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "        (k_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "        (v_proj): Linear(in_features=256, out_features=128, bias=True)\n",
              "        (out_proj): Linear(in_features=128, out_features=256, bias=True)\n",
              "      )\n",
              "      (norm_final_attn): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "    )\n",
              "    (iou_token): Embedding(1, 256)\n",
              "    (mask_tokens): Embedding(4, 256)\n",
              "    (unet_upscaling_layers): ModuleList(\n",
              "      (0): Up(\n",
              "        (up): ConvTranspose2dUsingLinear(\n",
              "          (upsampler_linear): Sequential(\n",
              "            (0): Linear(in_features=256, out_features=512, bias=True)\n",
              "            (1): ReLU()\n",
              "          )\n",
              "          (norm): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "        (conv): DoubleConv(\n",
              "          (double_conv): Sequential(\n",
              "            (0): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "            (1): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (1): Up(\n",
              "        (up): ConvTranspose2dUsingLinear(\n",
              "          (upsampler_linear): Sequential(\n",
              "            (0): Linear(in_features=256, out_features=512, bias=True)\n",
              "            (1): ReLU()\n",
              "          )\n",
              "          (norm): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "        (conv): DoubleConv(\n",
              "          (double_conv): Sequential(\n",
              "            (0): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(292, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "            (1): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(164, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (2): Up(\n",
              "        (up): ConvTranspose2dUsingLinear(\n",
              "          (upsampler_linear): Sequential(\n",
              "            (0): Linear(in_features=164, out_features=328, bias=True)\n",
              "            (1): ReLU()\n",
              "          )\n",
              "          (norm): BatchNorm2d(328, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "        (conv): DoubleConv(\n",
              "          (double_conv): Sequential(\n",
              "            (0): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(210, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "            (1): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (3-5): 3 x Up(\n",
              "        (up): ConvTranspose2dUsingLinear(\n",
              "          (upsampler_linear): Sequential(\n",
              "            (0): Linear(in_features=128, out_features=256, bias=True)\n",
              "            (1): ReLU()\n",
              "          )\n",
              "          (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "        (conv): DoubleConv(\n",
              "          (double_conv): Sequential(\n",
              "            (0): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(192, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "            (1): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "      (6): Up(\n",
              "        (up): ConvTranspose2dUsingLinear(\n",
              "          (upsampler_linear): Sequential(\n",
              "            (0): Linear(in_features=128, out_features=256, bias=True)\n",
              "            (1): ReLU()\n",
              "          )\n",
              "          (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "        (conv): DoubleConv(\n",
              "          (double_conv): Sequential(\n",
              "            (0): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "            (1): Conv2dWithActivation(\n",
              "              (conv): Sequential(\n",
              "                (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "                (1): ReLU()\n",
              "              )\n",
              "              (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "              (quant_output_normalization): QuantStub()\n",
              "            )\n",
              "          )\n",
              "        )\n",
              "      )\n",
              "    )\n",
              "    (output_hypernetworks_mlps): ModuleList(\n",
              "      (0-3): 4 x MLPBlock(\n",
              "        (layers): ModuleList(\n",
              "          (0-1): 2 x Sequential(\n",
              "            (0): Linear(in_features=256, out_features=256, bias=True)\n",
              "            (1): ReLU()\n",
              "          )\n",
              "        )\n",
              "        (fc): Linear(in_features=256, out_features=64, bias=True)\n",
              "      )\n",
              "    )\n",
              "    (iou_prediction_head): MLPBlock(\n",
              "      (layers): ModuleList(\n",
              "        (0-1): 2 x Sequential(\n",
              "          (0): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (1): ReLU()\n",
              "        )\n",
              "      )\n",
              "      (fc): Linear(in_features=256, out_features=4, bias=True)\n",
              "    )\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ],
      "source": [
        "from efficient_sam.build_efficient_sam import build_efficient_sam_vitt, build_efficient_sam_vits\n",
        "from squeeze_sam.build_squeeze_sam import build_squeeze_sam\n",
        "import zipfile\n",
        "\n",
        "efficient_sam_vitt_model = build_efficient_sam_vitt()\n",
        "efficient_sam_vitt_model.eval()\n",
        "\n",
        "# Since EfficientSAM-S checkpoint file is >100MB, we store the zip file.\n",
        "with zipfile.ZipFile(\"weights/efficient_sam_vits.pt.zip\", 'r') as zip_ref:\n",
        "    zip_ref.extractall(\"weights\")\n",
        "efficient_sam_vits_model = build_efficient_sam_vits()\n",
        "efficient_sam_vits_model.eval()\n",
        "\n",
        "squeeze_sam_model = build_squeeze_sam()\n",
        "squeeze_sam_model.eval()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b76-PTdKuidf"
      },
      "source": [
        "## Box segmentation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o6ed2uLDCSDn"
      },
      "source": [
        "prepare your own image here"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 195
        },
        "id": "Xq7KloryJrhf",
        "outputId": "2df80759-fbfe-48da-a11f-9ac3b9926ee6"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 3000x3000 with 4 Axes>"
            ],
            "image/png": "[encoded data, 466756 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "x1=400\n",
        "y1=200\n",
        "x2=800\n",
        "y2=600\n",
        "w=x2-x1\n",
        "h=y2-y1\n",
        "\n",
        "fig, ax = plt.subplots(1, 4, figsize=(30, 30))\n",
        "input_point = np.array([[x1, y1], [x2, y2]])\n",
        "input_label = np.array([2,3])\n",
        "image_path = \"figs/examples/dogs.jpg\"\n",
        "image = np.array(Image.open(image_path))\n",
        "show_points(input_point, input_label, ax[0])\n",
        "show_box([x1,y1,x2,y2], ax[0])\n",
        "ax[0].imshow(image)\n",
        "\n",
        "\n",
        "ax[1].imshow(image)\n",
        "mask_efficient_sam_vitt = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vitt_model)\n",
        "show_anns_ours(mask_efficient_sam_vitt, ax[1])\n",
        "ax[1].title.set_text(\"EfficientSAM (VIT-tiny)\")\n",
        "ax[1].axis('off')\n",
        "\n",
        "ax[2].imshow(image)\n",
        "mask_efficient_sam_vits = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vits_model)\n",
        "show_anns_ours(mask_efficient_sam_vits, ax[2])\n",
        "ax[2].title.set_text(\"EfficientSAM (VIT-small)\")\n",
        "ax[2].axis('off')\n",
        "\n",
        "\n",
        "ax[3].imshow(image)\n",
        "mask_squeeze_sam = run_ours_box_or_points(image_path, input_point, input_label, squeeze_sam_model)\n",
        "show_anns_ours(mask_squeeze_sam, ax[3])\n",
        "ax[3].title.set_text(\"SqueezeSAM\")\n",
        "ax[3].axis('off')\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6IQBINppEQXW"
      },
      "source": [
        "## Point segmentation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "9vpU0W2ewjYv",
        "outputId": "e2ff6fbe-303d-4f22-96b2-e7a510498373",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 195
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 3000x3000 with 4 Axes>"
            ],
            "image/png": "[encoded data, 477196 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "fig, ax = plt.subplots(1, 4, figsize=(30, 30))\n",
        "\n",
        "input_label = np.array([1, 1])\n",
        "image_path = \"figs/examples/dogs.jpg\"\n",
        "input_point = np.array([[580, 350], [650, 350]])\n",
        "image = np.array(Image.open(image_path))\n",
        "show_points(input_point, input_label, ax[0])\n",
        "ax[0].imshow(image)\n",
        "\n",
        "\n",
        "ax[1].imshow(image)\n",
        "mask_efficient_sam_vitt = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vitt_model)\n",
        "show_anns_ours(mask_efficient_sam_vitt, ax[1])\n",
        "ax[1].title.set_text(\"EfficientSAM (VIT-tiny)\")\n",
        "ax[1].axis('off')\n",
        "\n",
        "ax[2].imshow(image)\n",
        "mask_efficient_sam_vits = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vits_model)\n",
        "show_anns_ours(mask_efficient_sam_vits, ax[2])\n",
        "ax[2].title.set_text(\"EfficientSAM (VIT-small)\")\n",
        "ax[2].axis('off')\n",
        "\n",
        "\n",
        "ax[3].imshow(image)\n",
        "mask_squeeze_sam = run_ours_box_or_points(image_path, input_point, input_label, squeeze_sam_model)\n",
        "show_anns_ours(mask_squeeze_sam, ax[3])\n",
        "ax[3].title.set_text(\"SqueezeSAM\")\n",
        "ax[3].axis('off')\n",
        "\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "custom": {
      "cells": [],
      "metadata": {
        "custom": {
          "cells": [],
          "metadata": {
            "accelerator": "GPU",
            "colab": {
              "machine_shape": "hm",
              "provenance": []
            },
            "fileHeader": "",
            "fileUid": "f337ddbb-4ec7-4bc4-8c8b-f31305249752",
            "isAdHoc": false,
            "kernelspec": {
              "display_name": "Python 3",
              "name": "python3"
            },
            "language_info": {
              "name": "python"
            }
          },
          "nbformat": 4,
          "nbformat_minor": 0
        },
        "fileHeader": "",
        "fileUid": "e9a56628-4146-43e5-85f7-3d332cf3b1a2",
        "indentAmount": 2,
        "isAdHoc": false,
        "language_info": {
          "name": "plaintext"
        }
      },
      "nbformat": 4,
      "nbformat_minor": 2
    },
    "indentAmount": 2,
    "colab": {
      "provenance": []
    },
    "language_info": {
      "name": "python"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}