notebooks/community/model_garden/model_garden_pytorch_sam.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ad30fe2-1fc1-47e3-8a9f-624170b5aae6"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cYYdTBQoKeCP"
      },
      "source": [
        " # Vertex AI Model Garden - Segment Anything Model (SAM) Serving on Vertex AI\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "Open in Vertex AI Workbench\n",
        "    </a>\n",
        "    (a Python-3 CPU notebook is recommended)\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JbmPgTp2LRCY"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates using the [huggingface/transformers](https://github.com/huggingface/transformers) framework to serve Segment Anything Model (SAM) models and deploy them for online prediction on Vertex AI.\n",
        "\n",
        "Following the notebook you will conduct experiments using the pre-built docker image on Vertex AI.\n",
        "\n",
        "- With the pre-built docker images, you can **deploy** models for the following tasks:\n",
        "    - Mask Generation\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for image captioning.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "78f72e0a-52e5-4de5-ac0f-2171b3493825"
      },
      "source": [
        "## Setup environment\n",
        "\n",
        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G__krby2Mqmh"
      },
      "source": [
        "### Colab only"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RS5_6QVFMyc-"
      },
      "outputs": [],
      "source": [
        "!pip3 install --upgrade google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Li2aj1n9NDYl"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth as google_auth\n",
        "\n",
        "google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6c1bc20-3495-448a-b242-01930ba8153c"
      },
      "source": [
        "### Setup Google Cloud project\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
        "\n",
        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5b1f2c08-c84e-4158-a976-e6b59d1055ca"
      },
      "source": [
        "It's highly recommended to run this notebook on [Vertex AI workbench](https://cloud.google.com/vertex-ai-workbench)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fc1fc14-2d77-4bf7-8f6d-c1afc10c848a"
      },
      "source": [
        "If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk) and [gsutil](https://cloud.google.com/storage/docs/gsutil_install)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e02b9811-b730-4573-83fa-9d47f2ce0436"
      },
      "source": [
        "####Fill following variables for experiments environment:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7c2ce2fa-5a9b-40f6-b99d-6c1325775b36"
      },
      "outputs": [],
      "source": [
        "# Cloud project id.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The region you want to launch jobs in.\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
        "GCS_BUCKET = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fwgeVkFXOS64"
      },
      "source": [
        "*Initialize* Vertex AI API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MxqnJMahPK8r"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tSjMVyG8PYS-"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H05BzPO5Pcnh"
      },
      "outputs": [],
      "source": [
        "# The pre-built serving docker image.\n",
        "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
        "SERVE_DOCKER_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/sam-serve\"\n",
        ")\n",
        "\n",
        "# The serving port.\n",
        "SERVE_PORT = 7080\n",
        "\n",
        "# The serving route.\n",
        "SERVE_ROUTE = \"/predictions/sam_serving\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "82058975-23b5-4b97-9e14-dd9a29c578ed"
      },
      "source": [
        "### Define common utility functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a78f988d-e5e2-4a57-ba0f-569e970514c0"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "from io import BytesIO\n",
        "\n",
        "import matplotlib.patches as patches\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pycocotools.mask as mask_util\n",
        "import requests\n",
        "from PIL import Image\n",
        "\n",
        "\n",
        "def download_image(url):\n",
        "    response = requests.get(url)\n",
        "    return Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
        "\n",
        "\n",
        "def image_to_base64(image):\n",
        "    buffer = BytesIO()\n",
        "    image.save(buffer, format=\"JPEG\")\n",
        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
        "    return image_str\n",
        "\n",
        "\n",
        "def base64_to_image(image_str):\n",
        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
        "    return image\n",
        "\n",
        "\n",
        "def image_grid(imgs, rows=2, cols=2):\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
        "    return grid\n",
        "\n",
        "\n",
        "def draw_image_with_boxes(image, boxes):\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.imshow(image)\n",
        "    for box in boxes:\n",
        "        x, y = box[\"xmin\"], box[\"ymin\"]\n",
        "        width, height = box[\"xmax\"] - x, box[\"ymax\"] - y\n",
        "        rect = patches.Rectangle(\n",
        "            (x, y), width, height, linewidth=2, edgecolor=\"yellow\", facecolor=\"none\"\n",
        "        )\n",
        "        ax.add_patch(rect)\n",
        "    plt.axis(\"off\")\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def decode_rle_masks(pred_masks_rle):\n",
        "    return np.stack([mask_util.decode(rle) for rle in pred_masks_rle])\n",
        "\n",
        "\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([30 / 255, 144 / 255, 255 / 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",
        "\n",
        "def show_predictions(preds):\n",
        "    fig = plt.figure(figsize=(10, 7))\n",
        "\n",
        "    fig.add_subplot(1, 2, 1)\n",
        "    plt.imshow(np.array(image1))\n",
        "    ax = plt.gca()\n",
        "    masks = decode_rle_masks(preds[0][\"masks_rle\"])\n",
        "    for mask in masks:\n",
        "        show_mask(mask, ax=ax, random_color=True)\n",
        "    plt.axis(\"off\")\n",
        "\n",
        "    fig.add_subplot(1, 2, 2)\n",
        "    plt.imshow(np.array(image2))\n",
        "    ax = plt.gca()\n",
        "    masks = decode_rle_masks(preds[1][\"masks_rle\"])\n",
        "    for mask in masks:\n",
        "        show_mask(mask, ax=ax, random_color=True)\n",
        "    plt.axis(\"off\")\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def deploy_model(task, model_id):\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{task}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=task,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[SERVE_PORT],\n",
        "        serving_container_predict_route=SERVE_ROUTE,\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=\"n1-standard-8\",\n",
        "        accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "        accelerator_count=1,\n",
        "        deploy_request_timeout=1800,\n",
        "    )\n",
        "    return endpoint, model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "242fc8cf-5a0e-483b-8c9d-77a474cedc4b"
      },
      "source": [
        "## Mask Generation with Vertex AI Endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ot2HhTqxRYri"
      },
      "outputs": [],
      "source": [
        "endpoint, model = deploy_model(\n",
        "    task=\"mask-generation\", model_id=\"facebook/sam-vit-large\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d6e51c57-b5e2-4ae7-888a-5391cceee5fb"
      },
      "outputs": [],
      "source": [
        "image1 = download_image(\"http://images.cocodataset.org/val2017/000000039769.jpg\")\n",
        "image2 = download_image(\"http://images.cocodataset.org/val2017/000000000285.jpg\")\n",
        "grid = image_grid([image1, image2], 1, 2)\n",
        "display(grid)\n",
        "\n",
        "instances = [\n",
        "    {\"image\": image_to_base64(image1)},\n",
        "    {\"image\": image_to_base64(image2)},\n",
        "]\n",
        "preds = endpoint.predict(instances=instances).predictions\n",
        "show_predictions(preds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5UdquXMWR1E4"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TGTWIJK8R136"
      },
      "outputs": [],
      "source": [
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_sam.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}