notebooks/community/model_garden/model_garden_pytorch_owlvit.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "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": "2bd716bf3e39"
      },
      "source": [
        "# Vertex AI Model Garden - OWL-ViT\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_owlvit.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_owlvit.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_owlvit.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": "d8cd12648da4"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying the pre-trained [OWL-ViT](https://huggingface.co/google/owlvit-base-patch32) model on Vertex AI for online prediction.\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": "264c07757582"
      },
      "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": "d73ffa0c0b83"
      },
      "source": [
        "### Colab only"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "if \"google.colab\" in str(get_ipython()):\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()\n",
        "\n",
        "    # Restart the notebook kernel after installs.\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0f826ff482a2"
      },
      "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": "8958ebc71868"
      },
      "source": [
        "Fill following variables for experiments environment:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9db30f827a65"
      },
      "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": "92f16e22c20b"
      },
      "source": [
        "Initialize Vertex AI API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1680c257acfb"
      },
      "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": "6ca48b699d17"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "de9882ea89ea"
      },
      "outputs": [],
      "source": [
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10188266a5cd"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cac4478ae098"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import os\n",
        "from datetime import datetime\n",
        "from io import BytesIO\n",
        "\n",
        "import matplotlib.patches as patches\n",
        "import matplotlib.pyplot as plt\n",
        "import requests\n",
        "from google.cloud import aiplatform\n",
        "from PIL import Image\n",
        "\n",
        "\n",
        "def create_job_name(prefix):\n",
        "    user = os.environ.get(\"USER\")\n",
        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
        "    job_name = f\"{prefix}-{user}-{now}\"\n",
        "    return job_name\n",
        "\n",
        "\n",
        "def download_image(url):\n",
        "    response = requests.get(url)\n",
        "    return Image.open(BytesIO(response.content))\n",
        "\n",
        "\n",
        "def image_to_base64(image, format=\"JPEG\"):\n",
        "    buffer = BytesIO()\n",
        "    image.save(buffer, format=format)\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",
        "    plt.axis(\"off\")\n",
        "    ax.imshow(image)\n",
        "    if len(boxes) == 0:\n",
        "        return\n",
        "    boxes = boxes[\"boxes\"]\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.show()\n",
        "\n",
        "\n",
        "def deploy_model(model_id, task):\n",
        "    model_name = \"owl-vit\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
        "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/transformers_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "        artifact_uri=artifact_uri,\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=\"n1-standard-8\",\n",
        "        accelerator_type=\"NVIDIA_TESLA_T4\",\n",
        "        accelerator_count=1,\n",
        "        deploy_request_timeout=1800,\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2d72ecdb8c9"
      },
      "source": [
        "## Upload and deploy models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9448c5f545fa"
      },
      "source": [
        "This section uploads the pre-trained model to Model Registry and deploys it on the Endpoint with 1 T4 GPU.\n",
        "\n",
        "The model deployment step will take ~15 minutes to complete.\n",
        "\n",
        "Once deployed, you can send images and object texts to get bounding boxes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b4b46c28d8b1"
      },
      "outputs": [],
      "source": [
        "model, endpoint = deploy_model(\n",
        "    model_id=\"google/owlvit-base-patch32\", task=\"zero-shot-object-detection\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6be655247cb1"
      },
      "outputs": [],
      "source": [
        "image = download_image(\"http://images.cocodataset.org/val2017/000000039769.jpg\")\n",
        "\n",
        "instances = [\n",
        "    {\"image\": image_to_base64(image), \"text\": \"cat\"},\n",
        "]\n",
        "preds = endpoint.predict(instances=instances).predictions\n",
        "draw_image_with_boxes(image, preds[0])\n",
        "print(preds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "db7ffebdb4be"
      },
      "source": [
        "### Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ccf3714dbe9"
      },
      "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_owlvit.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}