Segment_Anything_Benchmarks.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "V100",
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Segment Anything Keras 3.0 port Benchmarks\n",
        "\n",
        "This notebook benchmarks the segment anything model for TensorFlow, JAX, and PyTorch using Keras 3.0.\n",
        "\n",
        "There are three types of benchmarks:\n",
        "\n",
        "1. End-to-end model inference (`image_encoder + prompt_encoder + mask_decoder`)\n",
        "2. End-to-end model inference with pre and post-processing\n",
        "3. Prompt benchmarks (`prompt_encoder + mask_decoder` with image features set)"
      ],
      "metadata": {
        "id": "belFxJbxUSuK"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Get all the dependencies and weight sets"
      ],
      "metadata": {
        "id": "3nS0Cb8hU3eV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "wvTf8fxHSp5c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7c1c8dbe-dc5a-4e4b-fac4-6e09fdb2b569"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "tensorflow 2.15.0 requires keras<2.16,>=2.15.0, but you have keras 3.0.2 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "# Get the dependencies\n",
        "!pip install -Uq keras-cv >> /dev/null\n",
        "!pip install -Uq keras >> /dev/null\n",
        "!pip install -Uq git+https://github.com/tirthasheshpatel/segment_anything_keras.git >> /dev/null\n",
        "\n",
        "# Get the image for the demo\n",
        "!curl -sSL https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg -o truck.jpg\n",
        "!curl -sSL https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/groceries.jpg -o groceries.jpg"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set the backend"
      ],
      "metadata": {
        "id": "wdSmBxslU8Xd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.environ['KERAS_BACKEND'] = \"torch\""
      ],
      "metadata": {
        "id": "fXasdeUnS0pb"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Choose the model"
      ],
      "metadata": {
        "id": "IfNiHXYLhsvQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_type = \"huge\""
      ],
      "metadata": {
        "id": "xykkspqNhuvS"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Import Dependencies"
      ],
      "metadata": {
        "id": "vSxB0WbKVaGw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import keras\n",
        "from keras import ops\n",
        "from keras_cv.models import SegmentAnythingModel\n",
        "from sam_keras import SAMPredictor"
      ],
      "metadata": {
        "id": "XHiX11_QS0rr"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Define the model"
      ],
      "metadata": {
        "id": "YMKf6hiLVDqf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sam = SegmentAnythingModel.from_preset(f\"sam_{model_type}_sa1b\")"
      ],
      "metadata": {
        "id": "38UUq9vzS0uY",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "eb01e05f-3272-4c2d-f4ba-1a4daa754365"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from https://storage.googleapis.com/keras-cv/models/segment_anything/sam_huge.h5\n",
            "\u001b[1m2564774344/2564774344\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m74s\u001b[0m 0us/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## End-to-End Model Inference with pre and post-processing"
      ],
      "metadata": {
        "id": "i1OZ0AaOVOul"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup"
      ],
      "metadata": {
        "id": "tXn-pElWhQdw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define predictor\n",
        "model = SAMPredictor(sam)\n",
        "transform  = model.transform\n",
        "\n",
        "# Load the image\n",
        "image = cv2.imread('truck.jpg')\n",
        "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "# Define the inputs\n",
        "input_point = np.array([[500, 375]])\n",
        "input_label = np.array([1])\n",
        "\n",
        "image_record = {}\n",
        "\n",
        "image_record[\"image\"] = ops.convert_to_tensor(\n",
        "    transform.apply_image(image)[np.newaxis, ...],\n",
        "    dtype=\"float32\"\n",
        ")\n",
        "\n",
        "image_record[\"original_size\"] = (image.shape[0], image.shape[1])\n",
        "\n",
        "image_record[\"point_coords\"] = ops.reshape(\n",
        "    ops.convert_to_tensor(\n",
        "        input_point, dtype=\"float32\"\n",
        "    ),\n",
        "    (1, 1, 2)\n",
        ")\n",
        "image_record[\"point_coords\"] = transform.apply_coords(\n",
        "    image_record[\"point_coords\"], image_record[\"original_size\"]\n",
        ")\n",
        "\n",
        "image_record[\"point_labels\"] = ops.convert_to_tensor(\n",
        "    input_label[np.newaxis, ...],\n",
        "    dtype=\"float32\"\n",
        ")"
      ],
      "metadata": {
        "id": "Ps93S2vVTZie"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Benchmark"
      ],
      "metadata": {
        "id": "UDqGr2i7hT-7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Dry run to build the model\n",
        "out = model.predict(image_record)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "N0CmG5a5TZk8",
        "outputId": "9fd3b98d-c5cd-4d24-8175-46c5d89f1b26"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m3s\u001b[0m 3s/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Predict also reports a time. Let's consider that too.\n",
        "out = model.predict(image_record)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6N_xMBuhRhtu",
        "outputId": "e1218f55-fd33-4888-e373-729751890ef4"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 598ms/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Benchmark the model\n",
        "%timeit out = model.predict(image_record, verbose=0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "We_RwF7dWY63",
        "outputId": "b16161c9-6c03-4db6-a939-fe4cbdcc5753"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "609 ms ± 3.58 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## End-to-End Model Inference"
      ],
      "metadata": {
        "id": "AKE1PdWqYLYJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup"
      ],
      "metadata": {
        "id": "Q4rm45Y7hXhF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Run the pre and post-processing steps here itself\n",
        "images = model.preprocess_images(image_record[\"image\"])\n",
        "points = ops.convert_to_tensor(\n",
        "    image_record.get(\"point_coords\", ops.ones((1, 0, 2))),\n",
        "    dtype=\"float32\"\n",
        ")\n",
        "labels = ops.convert_to_tensor(\n",
        "    image_record.get(\"point_labels\", ops.ones((1, 0))),\n",
        "    dtype=\"float32\"\n",
        ")\n",
        "box = ops.convert_to_tensor(\n",
        "    image_record.get(\"boxes\", ops.ones((1, 0, 2, 2))),\n",
        "    dtype=\"float32\"\n",
        ")\n",
        "mask = ops.convert_to_tensor(\n",
        "    image_record.get(\"mask_inputs\", ops.ones((1, 0, 256, 256, 1))),\n",
        "    dtype=\"float32\"\n",
        ")\n",
        "\n",
        "if ops.size(points) and not ops.size(box):\n",
        "    pad_point = ops.zeros((points.shape[0], 1, 2), dtype=\"float32\")\n",
        "    pad_label = -ops.ones((labels.shape[0], 1), dtype=\"float32\")\n",
        "    points = ops.concatenate([points, pad_point], axis=1)\n",
        "    labels = ops.concatenate([labels, pad_label], axis=1)\n",
        "\n",
        "B = max([\n",
        "    images.shape[0],\n",
        "    points.shape[0],\n",
        "    labels.shape[0],\n",
        "    box.shape[0],\n",
        "    mask.shape[0],\n",
        "])\n",
        "\n",
        "images, points, labels, box, mask = model._broadcast_batch(\n",
        "    B, images, points, labels, box, mask\n",
        ")\n",
        "\n",
        "model_input = {\n",
        "    \"images\": images,\n",
        "    \"points\": points,\n",
        "    \"labels\": labels,\n",
        "    \"boxes\": box,\n",
        "    \"masks\": mask\n",
        "}"
      ],
      "metadata": {
        "id": "835jlxIyXkvv"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Benchmark"
      ],
      "metadata": {
        "id": "DoovDjh5hbax"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.model.predict(model_input);"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jprz_7nqXJ2V",
        "outputId": "fd3be700-ac13-45b8-a3dc-4dc45d837ab3"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 600ms/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit model.model.predict(model_input, verbose=0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U27ga7D8YO12",
        "outputId": "5f70ff29-715f-4061-e159-46d6046ab19d"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "608 ms ± 2.13 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Prompt Benchmarks"
      ],
      "metadata": {
        "id": "nvaEd-AlaCf2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup"
      ],
      "metadata": {
        "id": "h_mCbABphgGC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Set the features\n",
        "features = ops.convert_to_tensor(\n",
        "    model.model.backbone.predict(model_input[\"images\"], verbose=0),\n",
        "    dtype=\"float32\"\n",
        ")"
      ],
      "metadata": {
        "id": "eKgwlNw7YTem"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class SAMPrompter(keras.Model):\n",
        "    def __init__(self, prompt_encoder, mask_decoder, feature_shape=(64, 64, 256), **kwargs):\n",
        "        # Define the prompt encoder inputs -- Prompts\n",
        "        prompt_inputs = {\n",
        "            \"points\": keras.Input(shape=[None, 2], name=\"points\"),\n",
        "            \"labels\": keras.Input(shape=[None], name=\"labels\"),\n",
        "            \"boxes\": keras.Input(shape=[None, 2, 2], name=\"boxes\"),\n",
        "            \"masks\": keras.Input(shape=[None, None, None, 1], name=\"masks\"),\n",
        "        }\n",
        "\n",
        "        # All Inputs -- Features + Prompts\n",
        "        all_inputs = {\"features\": keras.Input(feature_shape, name=\"features\")}\n",
        "        all_inputs.update(prompt_inputs)\n",
        "\n",
        "        # Build the prompt encoder\n",
        "        prompt_embeddings = prompt_encoder(prompt_inputs)\n",
        "\n",
        "        # Define the mask decoder inputs\n",
        "        mask_decoder_inputs = {\n",
        "            \"image_embeddings\": all_inputs[\"features\"],\n",
        "            \"image_pe\": prompt_embeddings[\"dense_positional_embeddings\"],\n",
        "            \"sparse_prompt_embeddings\": prompt_embeddings[\"sparse_embeddings\"],\n",
        "            \"dense_prompt_embeddings\": prompt_embeddings[\"dense_embeddings\"],\n",
        "        }\n",
        "\n",
        "        # Build the mask decoder\n",
        "        outputs = mask_decoder(mask_decoder_inputs)\n",
        "\n",
        "        super().__init__(inputs=all_inputs, outputs=outputs, **kwargs)\n",
        "\n",
        "        self.prompt_encoder = prompt_encoder\n",
        "        self.mask_decoder = mask_decoder"
      ],
      "metadata": {
        "id": "SZif2w3rZSJs"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompter_model = SAMPrompter(model.model.prompt_encoder, model.model.mask_decoder, feature_shape=features.shape[1:])"
      ],
      "metadata": {
        "id": "9xEISA_IgKa9"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompt_inputs = {\n",
        "    \"features\": features,\n",
        "    \"points\": model_input[\"points\"],\n",
        "    \"labels\": model_input[\"labels\"],\n",
        "    \"boxes\": model_input[\"boxes\"],\n",
        "    \"masks\": model_input[\"masks\"]\n",
        "}"
      ],
      "metadata": {
        "id": "UtwBbwv4gUa5"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Benchmark"
      ],
      "metadata": {
        "id": "nLRbgdeDhiyP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Dry run to build the model\n",
        "outs = prompter_model.predict(prompt_inputs)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ptDHJOtLgnuL",
        "outputId": "902c2cd4-62bd-4b42-e896-298c195ff601"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 38ms/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Predict also reports a time. Let's also consider that.\n",
        "outs = prompter_model.predict(prompt_inputs)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DBkGykXKRykq",
        "outputId": "7bde8332-3572-4062-cbbf-49c942e37c2e"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 38ms/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit outs = prompter_model.predict(prompt_inputs, verbose=0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9AGFkbYQg1Op",
        "outputId": "b0fe8765-a429-4d2d-fa96-15c4dc3a0021"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "39.4 ms ± 834 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "HbuSvHkQWGdc"
      },
      "execution_count": 19,
      "outputs": []
    }
  ]
}