notebooks/community/model_garden/model_garden_jax_fvlm.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - JAX F-VLM\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_jax_fvlm.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",
        "\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_fvlm.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>                                                                                               <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_jax_fvlm.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",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24743cf4a1e1"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates training and serving a [JAX F-VLM model](https://github.com/google-research/google-research/tree/master/fvlm) for [open-vocabulary object detection and instance segmentation](https://arxiv.org/abs/2209.15639) tasks and deploying them on Vertex AI for online prediction."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to:\n",
        "\n",
        "- Prepare a training dataset.\n",
        "- Train a new JAX F-VLM model.\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 open-vocabulary image object detection and instance segmentation.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Model Registry\n",
        "- Vertex AI Online Prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08d289fa873f"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This notebook uses the following prediction image as an example:\n",
        "\n",
        "Image: https://pixabay.com/nl/photos/het-fruit-eten-citroen-limoen-3134631/\n",
        "\n",
        "Creative Commons License: https://pixabay.com/nl/service/terms/\n",
        "\n",
        "You can use your own custom prediction image as well as by modifying the `DEMO_IMAGE_PATH` variable in this notebook below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "source": [
        "### 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": "i7EUnXsZhAGF"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "# Install the packages.\n",
        "! pip3 install --upgrade google-cloud-aiplatform\n",
        "! sudo apt-get install -y subversion protobuf-compiler python3-lxml\\\n",
        "  python3-pip python3-dev git unzip\n",
        "# Note: The following libraries are pinned down versions of:\n",
        "# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt\n",
        "! pip install tensorflow==2.12.0\n",
        "! pip install tensorflow-datasets==4.9.2\n",
        "! pip install numpy==1.23.5\n",
        "! pip install torch==2.0.1+cu118 --extra-index-url https://download.pytorch.org/whl/cu118\n",
        "! pip install torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118\n",
        "! pip install opencv-python==4.7.0.72\n",
        "! pip install tqdm==4.65.0\n",
        "! pip install git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33\n",
        "! pip install Pillow==9.1.1\n",
        "! pip install orbax-checkpoint==0.3.3\n",
        "! pip install gin-config==0.5.0\n",
        "! pip install pycocotools==2.0.6\n",
        "! pip install contextlib2==21.6.0\n",
        "! pip install ml-collections==0.1.1\n",
        "! pip install chex==0.1.7\n",
        "! pip install optax==0.1.5\n",
        "# Dependencies already included. Use no-deps to not update numpy.\n",
        "! pip install --no-deps flax==0.7.2\n",
        "! pip install --no-deps clu==0.0.9\n",
        "! pip install jax[cuda11_cudnn86]==0.4.9 \\\n",
        "  --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
        "! pip install Cython\n",
        "! pip install git+https://github.com/cocodataset/cocoapi#subdirectory=PythonAPI\n",
        "\n",
        "# Get F-VLM repository by using svn to avoid downloading entire google-research repository.\n",
        "%cd\n",
        "! rm -rf ./fvlm\n",
        "! svn export -r 60422 https://github.com/google-research/google-research/trunk/fvlm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_9q83As4G2Yn"
      },
      "source": [
        "Download the F-VLM checkpoints into the `fvlm/checkpoints` folder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eXVI6s57FPyg"
      },
      "outputs": [],
      "source": [
        "%cd fvlm/checkpoints\n",
        "! ./download.sh\n",
        "%cd ../../"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7237befefbb2"
      },
      "source": [
        "Download COCO embeddings and setup data preparation code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "324b37791d24"
      },
      "outputs": [],
      "source": [
        "%cd fvlm\n",
        "! bash ./scripts/download_precomputed_embeddings.sh\n",
        "! git clone https://github.com/tensorflow/tpu.git\n",
        "! git clone http://github.com/tensorflow/models tf-models\n",
        "%cd tf-models/research\n",
        "! protoc object_detection/protos/*.proto --python_out=.\n",
        "%cd ../../../"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Restart kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "# Restart kernel after installs so that your environment can access the new packages.\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\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. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Check your accelerator [quota](https://console.cloud.google.com/iam-admin/quotas). This notebook uses TPU V3 8 cores, you can filter and request `Custom model training TPU V3 cores per region` quota for 8 cores for the `Vertex AI API` service in your region.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "twgKk-LsLmX3"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74ccc9e52986"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de775a3773ba"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6b2ccc891ed"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import datetime\n",
        "import functools\n",
        "import json\n",
        "import os\n",
        "import sys\n",
        "from io import BytesIO\n",
        "\n",
        "import gin\n",
        "import jax\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tqdm\n",
        "from PIL import Image\n",
        "\n",
        "# Work in fvlm directory in order to create gin config later.\n",
        "os.chdir(os.path.join(os.path.expanduser(\"~\"), \"fvlm\"))\n",
        "sys.path.append(os.getcwd())\n",
        "from demo_utils import input_utils as inputs\n",
        "from demo_utils import vis_utils\n",
        "from google.cloud import aiplatform\n",
        "from google.protobuf import json_format\n",
        "from google.protobuf.struct_pb2 import Value\n",
        "from utils import clip_utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vS1hQiGuLmX4"
      },
      "outputs": [],
      "source": [
        "staging_bucket = os.path.join(BUCKET_URI, \"jax_fvlm_staging\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=staging_bucket)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2cc825514deb"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b42bd4fa2b2d"
      },
      "outputs": [],
      "source": [
        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
        "\n",
        "# The F-VLM model to use. Choose between 'resnet_50', 'resnet_50x4', or 'resnet_50x16'.\n",
        "MODEL = \"resnet_50\"\n",
        "\n",
        "# COCO dataset constants. Only used for training with COCO dataset.\n",
        "BASE_ANNOTATION_URL = \"http://images.cocodataset.org/annotations\"\n",
        "INSTANCES_FILE = \"annotations_trainval2017.zip\"\n",
        "IMAGE_INFO_FILE = \"image_info_test2017.zip\"\n",
        "BASE_IMAGE_URL = \"http://images.cocodataset.org/zips\"\n",
        "TRAIN_IMAGE_FILE = \"train2017.zip\"\n",
        "VAL_IMAGE_FILE = \"val2017.zip\"\n",
        "TRAIN_ANNOTATION_PATH = \"./annotations/instances_train2017.json\"\n",
        "VAL_ANNOTATION_PATH = \"./annotations/instances_val2017.json\"\n",
        "LOCAL_CATEGORY_EMBEDDING_PATH = f\"./embeddings/{MODEL}/coco_embed.npy\"\n",
        "# Paths to save annotations for training example.\n",
        "TRAIN_TOY_ANNOTATION_PATH = \"./annotations/toy_instances_train2017.json\"\n",
        "VAL_TOY_ANNOTATION_PATH = \"./annotations/toy_instances_val2017.json\"\n",
        "TRAIN_IMAGE_DIR = \"./train2017\"\n",
        "VAL_IMAGE_DIR = \"./val2017\"\n",
        "TRAIN_DATA_PREFIX = \"train\"\n",
        "VAL_DATA_PREFIX = \"val\"\n",
        "# Local directory to save tfrecord dataset.\n",
        "LOCAL_TFRECORD_DIR = \"./tfrecord\"\n",
        "# GCS directory to upload tfrecord dataset.\n",
        "GCS_TFRECORD_PATH = f\"{BUCKET_URI}/dataset\"\n",
        "\n",
        "# GCS location of the category embeddings file.\n",
        "GCS_CATEGORY_EMBEDDING_PATH = f\"{BUCKET_URI}/embeddings/{MODEL}/embed.npy\"\n",
        "# An upper bound on the number of classes.\n",
        "MAX_NUM_CLS = 91\n",
        "# URL to download checkpoint.\n",
        "CHECKPOINT_URL = (\n",
        "    \"https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/fvlm/\"\n",
        "    f\"jax_checkpoints/{MODEL.replace('resnet_', 'r')}_checkpoint_184000\"\n",
        ")\n",
        "\n",
        "# The local path to the F-VLM folder.\n",
        "F_VLM_FOLDER = \".\"  # Current directory.\n",
        "# Train config template path.\n",
        "LOCAL_TRAIN_CONFIG_TEMPLATE_PATH = f\"{F_VLM_FOLDER}/configs/fvlm_train_and_eval.gin\"\n",
        "# Final train config path.\n",
        "GCS_TRAIN_CONFIG_PATH = f\"{BUCKET_URI}/fvlm_train_and_eval.gin\"\n",
        "# Base output path for training artifacts.\n",
        "GCS_TRAIN_BASE_PATH = f\"{BUCKET_URI}/train\"\n",
        "# Final train config path.\n",
        "GCS_TRAIN_CONFIG_PATH = f\"{GCS_TRAIN_BASE_PATH}/fvlm_train_and_eval.gin\"\n",
        "# Training container image.\n",
        "TRAIN_DOCKER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/jax-f-vlm-train\"\n",
        "GCS_TRAIN_DATA_PATH = f\"{GCS_TFRECORD_PATH}/{TRAIN_DATA_PREFIX}*.tfrecord\"\n",
        "GCS_VAL_DATA_PATH = f\"{GCS_TFRECORD_PATH}/{VAL_DATA_PREFIX}*.tfrecord\"\n",
        "GCS_VAL_ANNOTATION_FILE = f\"{GCS_TRAIN_BASE_PATH}/val_annotation.json\"\n",
        "LOCAL_CHECKPOINT_FILE = (\n",
        "    f\"./checkpoints/{MODEL.replace('resnet_', 'r')}_checkpoint_184000\"\n",
        ")\n",
        "# GCS location to upload pretrained checkpoint to.\n",
        "GCS_CHECKPOINT_DIR = f\"{GCS_TRAIN_BASE_PATH}/pretrained_checkpoint\"\n",
        "\n",
        "# The train/eval batch sizes should be divisible by the number of GPUs used.\n",
        "TRAIN_BATCH_SIZE = 8\n",
        "EVAL_BATCH_SIZE = 8\n",
        "# Whether to predict mask.\n",
        "# Set False for object detection task and True for segmentation task.\n",
        "INCLUDE_MASK = True\n",
        "# Number of TPU cores to be used.\n",
        "NUM_CORES = 8\n",
        "# Total train steps.\n",
        "TRAIN_STEPS = 1000  # 1 epoch is 3787 steps in this example.\n",
        "# Total evaluation steps.\n",
        "EVAL_STEPS = 156\n",
        "# Initial learning rate.\n",
        "INIT_LEARNING_RATE = 0.001\n",
        "# Training image size.\n",
        "OUTPUT_SIZE = 1024\n",
        "# Variable dtype. bfloat16 or float32 are supported.\n",
        "DTYPE = \"bfloat16\"\n",
        "\n",
        "# The pre-built TF SavedModel conversion docker image.\n",
        "MODEL_CONVERSION_DOCKER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/jax-f-vlm-model-conversion\"\n",
        "\n",
        "# The pre-built prediction docker image.\n",
        "OPTIMIZED_TF_RUNTIME_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest\"\n",
        ")\n",
        "# The max number of boxes to draw on the output image.\n",
        "MAX_BOXES_TO_DRAW = 25\n",
        "# The minimum score required to draw a detected object.\n",
        "MIN_SCORE_THRESH = 0.2  # @param {type:\"slider\", min:0, max:0.9, step:0.05}\n",
        "# The local path to the output image.\n",
        "OUTPUT_IMAGE_PATH = \"./output.jpg\"\n",
        "# The F-VLM SavedModel folder which takes image and text embeddings as inputs.\n",
        "GCS_SAVED_MODEL_DIR = f\"{BUCKET_URI}/saved_model\"\n",
        "# The converted SavedModel folder which takes jpeg bytes and text-embeddings bytes as inputs.\n",
        "LOCAL_CONVERTED_SAVED_MODEL_DIR = f\"{F_VLM_FOLDER}/converted_saved_model\"\n",
        "# The Cloud Storage location for the converted SavedModel.\n",
        "GCS_CONVERTED_SAVED_MODEL_DIR = f\"{BUCKET_URI}/fvlm_saved_model\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c250872074f"
      },
      "source": [
        "### Define common functions\n",
        "\n",
        "This section defines functions for:\n",
        "\n",
        "- Loading and converting input image into the required prediction format.\n",
        "- Visualization of detection outputs.\n",
        "- Getting GCS Fuse path.\n",
        "- Getting a job name with current time.\n",
        "- Save a subset of COCO annotation file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XcYUGwr-AJGY"
      },
      "outputs": [],
      "source": [
        "def convert_numpy_array_to_byte_string_via_tf_tensor(np_array):\n",
        "    \"\"\"Serializes a numpy array to tensor bytes.\"\"\"\n",
        "    tensor_array = tf.convert_to_tensor(np_array)\n",
        "    tensor_byte_string = tf.io.serialize_tensor(tensor_array)\n",
        "    return tensor_byte_string.numpy()\n",
        "\n",
        "\n",
        "def generate_text_embeddings(categories):\n",
        "    \"\"\"Generates text embeddings in numpy format from object categories.\"\"\"\n",
        "    clip_text_fn = clip_utils.get_clip_text_fn(MODEL)\n",
        "    class_clip_features = []\n",
        "    print(\"Computing custom category text embeddings.\")\n",
        "    for cls_name in tqdm.tqdm(categories, total=len(categories)):\n",
        "        cls_feat = clip_text_fn(cls_name)\n",
        "        class_clip_features.append(cls_feat)\n",
        "    text_embeddings = np.concatenate(class_clip_features, axis=0)\n",
        "    embed_path = (\n",
        "        f'{F_VLM_FOLDER}/data/{MODEL.replace(\"resnet_\", \"r\")}_bg_empty_embed.npy'\n",
        "    )\n",
        "    background_embedding, empty_embeddings = np.load(embed_path)\n",
        "    background_embedding = background_embedding[np.newaxis, Ellipsis]\n",
        "    empty_embeddings = empty_embeddings[np.newaxis, Ellipsis]\n",
        "    tile_empty_embeddings = np.tile(\n",
        "        empty_embeddings, (MAX_NUM_CLS - len(categories) - 1, 1)\n",
        "    )\n",
        "    # Concatenate 'background' and 'empty' embeddings.\n",
        "    text_embeddings = np.concatenate(\n",
        "        (background_embedding, text_embeddings, tile_empty_embeddings), axis=0\n",
        "    )\n",
        "    return text_embeddings\n",
        "\n",
        "\n",
        "def get_jpeg_bytes(local_image_path, new_width=-1):\n",
        "    \"\"\"Returns jpeg bytes given an image path and resizes if required.\"\"\"\n",
        "    image = Image.open(local_image_path)\n",
        "    if new_width <= 0:\n",
        "        new_image = image\n",
        "    else:\n",
        "        width, height = image.size\n",
        "        print(\"original input image size: \", width, \" , \", height)\n",
        "        new_height = int(height * new_width / width)\n",
        "        print(\"new input image size: \", new_width, \" , \", new_height)\n",
        "        new_image = image.resize((new_width, new_height))\n",
        "    buffered = BytesIO()\n",
        "    new_image.save(buffered, format=\"JPEG\")\n",
        "    return buffered.getvalue()\n",
        "\n",
        "\n",
        "def generate_prediction_output_image(\n",
        "    input_image_path, prediction_output, output_image_path, categories\n",
        "):\n",
        "    \"\"\"Generates prediction output image with detected objects and bounding boxes.\"\"\"\n",
        "    # Generate tensors from prediction outputs.\n",
        "    prediction_output_tensor = {}\n",
        "    for key, val in prediction_output.items():\n",
        "        prediction_output_tensor[key] = tf.expand_dims(\n",
        "            tf.convert_to_tensor(val), axis=0\n",
        "        )\n",
        "    prediction_output_tensor[\"num_detections\"] = tf.cast(\n",
        "        prediction_output_tensor[\"num_detections\"], tf.int32\n",
        "    )\n",
        "    # Generate image embeddings for the input image.\n",
        "    with open(input_image_path, \"rb\") as f:\n",
        "        np_image = np.array(Image.open(f))\n",
        "    parser_fn = inputs.get_maskrcnn_parser()\n",
        "    data = parser_fn({\"image\": np_image, \"source_id\": np.array([0])})\n",
        "    np_data = jax.tree_map(lambda x: x.numpy()[np.newaxis, Ellipsis], data)\n",
        "    image_embeddings = np_data.pop(\"images\")\n",
        "    labels = np_data.pop(\"labels\")\n",
        "    # Generate visualization.\n",
        "    print(\"Preparing visualization.\")\n",
        "    id_mapping = {(i + 1): c for i, c in enumerate(categories)}\n",
        "    id_mapping[0] = \"background\"\n",
        "    for k in range(len(categories) + 2, MAX_NUM_CLS):\n",
        "        id_mapping[k] = \"empty\"\n",
        "    category_index = inputs.get_category_index(id_mapping)\n",
        "    maskrcnn_visualizer_fn = functools.partial(\n",
        "        vis_utils.visualize_boxes_and_labels_on_image_array,\n",
        "        category_index=category_index,\n",
        "        use_normalized_coordinates=False,\n",
        "        max_boxes_to_draw=MAX_BOXES_TO_DRAW,\n",
        "        min_score_thresh=MIN_SCORE_THRESH,\n",
        "        skip_labels=False,\n",
        "    )\n",
        "    vis_image = vis_utils.visualize_instance_segmentations(\n",
        "        prediction_output_tensor,\n",
        "        image_embeddings,\n",
        "        labels[\"image_info\"],\n",
        "        maskrcnn_visualizer_fn,\n",
        "    )\n",
        "    pil_vis_image = Image.fromarray(vis_image, mode=\"RGB\")\n",
        "    pil_vis_image.save(output_image_path)\n",
        "    print(\"Completed saving the output image at: \", output_image_path)\n",
        "\n",
        "\n",
        "def gcs_fuse_path(path: str) -> str:\n",
        "    \"\"\"Try to convert path to gcsfuse path if it starts with gs:// else do not modify it.\"\"\"\n",
        "    path = path.strip()\n",
        "    if path.startswith(\"gs://\"):\n",
        "        return \"/gcs/\" + path[5:]\n",
        "    return path\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str):\n",
        "    \"\"\"Gets a job name by adding current time to prefix.\"\"\"\n",
        "    return prefix + datetime.datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def save_subset_annotation(input_annotation_path, output_annotation_path):\n",
        "    \"\"\"Saves a subset of COCO annotation json file with CCA 4.0 license.\"\"\"\n",
        "\n",
        "    with open(input_annotation_path) as f:\n",
        "        coco_json = json.load(f)\n",
        "\n",
        "    img_ids = set()\n",
        "    images = []\n",
        "    annotations = []\n",
        "\n",
        "    for img in coco_json[\"images\"]:\n",
        "        if img[\"license\"] in [4, 5]:  # CCA 4.0 license.\n",
        "            img_ids.add(img[\"id\"])\n",
        "            images.append(img)\n",
        "\n",
        "    for ann in coco_json[\"annotations\"]:\n",
        "        if ann[\"image_id\"] in img_ids:\n",
        "            annotations.append(ann)\n",
        "\n",
        "    new_json = {\n",
        "        \"info\": coco_json[\"info\"],\n",
        "        \"licenses\": coco_json[\"licenses\"],\n",
        "        \"images\": images,\n",
        "        \"annotations\": annotations,\n",
        "        \"categories\": coco_json[\"categories\"],\n",
        "    }\n",
        "\n",
        "    with open(output_annotation_path, \"w\") as f:\n",
        "        json.dump(new_json, f)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ayNrua2txk0B"
      },
      "source": [
        "## Train new models\n",
        "\n",
        "This section shows how to train new models:\n",
        "\n",
        "1. Convert input data to tfrecord format.\n",
        "2. Upload tfrecord files and category embeddings to GCS location.\n",
        "3. Create Custom training job to train a new JAX model.\n",
        "\n",
        "If you want to use the provided pretrained saved model, you can skip this section and go to the `Convert F-VLM SavedModel to support smaller input size` section."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2d45a55bf581"
      },
      "source": [
        "### Prepare input data for training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "98c56d2a5ba8"
      },
      "source": [
        "Download the F-VLM SavedModels and checkpoints into the `fvlm/checkpoints` folder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a3f2d8ec7473"
      },
      "outputs": [],
      "source": [
        "%cd checkpoints\n",
        "! ./download.sh\n",
        "! wget {CHECKPOINT_URL}\n",
        "%cd ../../"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9c947b28adec"
      },
      "source": [
        "Download and unzip COCO annotations and images. This can take about 25 minutes. This section is needed only if you are using COCO dataset for training. If you want to use your own dataset, then skip this section and prepare your own dataset in COCO format by saving your train/validation images in `TRAIN_IMAGE_DIR` and `VAL_IMAGE_DIR` respectively, and write train/validation COCO format annotations at `TRAIN_TOY_ANNOTATION_PATH`and `VAL_TOY_ANNOTATION_PATH` respectively."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5a1e3561df85"
      },
      "outputs": [],
      "source": [
        "# Download and unzip instances file.\n",
        "! wget -nd -c {BASE_ANNOTATION_URL}/{INSTANCES_FILE}\n",
        "! unzip -nq {INSTANCES_FILE}\n",
        "\n",
        "# Download and unzip image info file.\n",
        "! wget -nd -c {BASE_ANNOTATION_URL}/{IMAGE_INFO_FILE}\n",
        "! unzip -nq {IMAGE_INFO_FILE}\n",
        "\n",
        "# Download and unzip train images.\n",
        "! wget -nd -c {BASE_IMAGE_URL}/{TRAIN_IMAGE_FILE}\n",
        "! unzip -nq {TRAIN_IMAGE_FILE}\n",
        "\n",
        "# Download and unzip eval images.\n",
        "! wget -nd -c {BASE_IMAGE_URL}/{VAL_IMAGE_FILE}\n",
        "! unzip -nq {VAL_IMAGE_FILE}\n",
        "\n",
        "# Remove zip files.\n",
        "! rm {INSTANCES_FILE}\n",
        "! rm {IMAGE_INFO_FILE}\n",
        "! rm {TRAIN_IMAGE_FILE}\n",
        "! rm {VAL_IMAGE_FILE}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dfa1084bba1b"
      },
      "source": [
        "Create toy annotation files with a subset of COCO annotations. This section is needed only if you are using COCO dataset for training. If you want to use your own dataset, then skip this section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "82ddec0cfd8d"
      },
      "outputs": [],
      "source": [
        "save_subset_annotation(\n",
        "    input_annotation_path=TRAIN_ANNOTATION_PATH,\n",
        "    output_annotation_path=TRAIN_TOY_ANNOTATION_PATH,\n",
        ")\n",
        "save_subset_annotation(\n",
        "    input_annotation_path=VAL_ANNOTATION_PATH,\n",
        "    output_annotation_path=VAL_TOY_ANNOTATION_PATH,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5f8d691d78c2"
      },
      "source": [
        "Save the dataset as TF Record. This can take about 10 minutes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "77be933a6ac9"
      },
      "outputs": [],
      "source": [
        "! PYTHONPATH=\"tf-models:tf-models/research\" python3 ./tpu/tools/datasets/create_coco_tf_record.py \\\n",
        "--logtostderr \\\n",
        "--include_masks \\\n",
        "--image_dir={TRAIN_IMAGE_DIR} \\\n",
        "--object_annotations_file={TRAIN_TOY_ANNOTATION_PATH} \\\n",
        "--output_file_prefix={LOCAL_TFRECORD_DIR}/{TRAIN_DATA_PREFIX} \\\n",
        "--num_shards=256"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8c8d56cc2085"
      },
      "outputs": [],
      "source": [
        "! PYTHONPATH=\"tf-models:tf-models/research\" python3 ./tpu/tools/datasets/create_coco_tf_record.py \\\n",
        "--include_masks \\\n",
        "--image_dir={VAL_IMAGE_DIR} \\\n",
        "--object_annotations_file={VAL_TOY_ANNOTATION_PATH} \\\n",
        "--output_file_prefix={LOCAL_TFRECORD_DIR}/{VAL_DATA_PREFIX} \\\n",
        "--num_shards=256"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "990866f0a3cb"
      },
      "source": [
        "### Upload tfrecord files, pretrained checkpoint and category embeddings to GCS location"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e1482d2bf054"
      },
      "outputs": [],
      "source": [
        "! gsutil -m rm -R -f {GCS_TFRECORD_PATH}\n",
        "! gsutil -m cp -R {LOCAL_TFRECORD_DIR} {GCS_TFRECORD_PATH}\n",
        "! gsutil -m cp {VAL_TOY_ANNOTATION_PATH} {GCS_VAL_ANNOTATION_FILE}\n",
        "! gsutil -m cp {LOCAL_CHECKPOINT_FILE} {GCS_CHECKPOINT_DIR}/checkpoint_0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "34c8961224f9"
      },
      "source": [
        "Upload category embedding. If you want to use your own dataset, then skip this cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b2514916b9cb"
      },
      "outputs": [],
      "source": [
        "! gsutil -m cp {LOCAL_CATEGORY_EMBEDDING_PATH} {GCS_CATEGORY_EMBEDDING_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "605a02f752e5"
      },
      "source": [
        "If you want to use your own dataset, uncomment the cell below to create your own category embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8f442f2c25c1"
      },
      "outputs": [],
      "source": [
        "# CATEGORIES = [\n",
        "#     \"kiwi\",\n",
        "#     \"orange\",\n",
        "#     \"lemon\",\n",
        "#     \"blackberry\",\n",
        "#     \"pine cone\",\n",
        "#     \"red orange\",\n",
        "#     \"table\",\n",
        "#     \"spoon\",\n",
        "#     \"pine needles\",\n",
        "#     \"seed\",\n",
        "# ]\n",
        "# text_embeddings = generate_text_embeddings(categories=CATEGORIES)\n",
        "# with tf.io.gfile.GFile(GCS_CATEGORY_EMBEDDING_PATH, \"w\") as f:\n",
        "#     np.save(f, text_embeddings)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd1b3a5d5171"
      },
      "source": [
        "### Create Custom training job and train a new model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8a30f2fa7d34"
      },
      "source": [
        "Set training hyperparameters. To see the full list of customizable hyperparameters, see the [config template file](https://github.com/google-research/google-research/blob/e1f9fae637db06ba885217518cbbf7f4fa4b9d7b/fvlm/configs/fvlm_train_and_eval.gin)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b2f5dd9951ff"
      },
      "outputs": [],
      "source": [
        "assert TRAIN_BATCH_SIZE % NUM_CORES == 0, \"NUM_CORES must divide TRAIN_BATCH_SIZE.\"\n",
        "assert EVAL_BATCH_SIZE % NUM_CORES == 0, \"NUM_CORES must divide EVAL_BATCH_SIZE.\"\n",
        "config_overrides = [\n",
        "    \"evaluate.host_evaluator = None\",  # For evaluation without COCO annotation json file.\n",
        "    \"evaluate.eval_metrics = {'coco_metric': @COCODetectionMetric}\",  # For evaluation without COCO annotation json file.\n",
        "    f\"TRAIN_FILE_PATTERN = '{gcs_fuse_path(GCS_TRAIN_DATA_PATH)}'\",\n",
        "    f\"EVAL_FILE_PATTERN = '{gcs_fuse_path(GCS_VAL_DATA_PATH)}'\",\n",
        "    f\"TRAIN_BS = {TRAIN_BATCH_SIZE}\",\n",
        "    f\"EVAL_BS = {EVAL_BATCH_SIZE}\",\n",
        "    f\"TRAIN_STEPS = {TRAIN_STEPS}\",\n",
        "    f\"EVAL_STEPS = {EVAL_STEPS}\",\n",
        "    f\"EMBED_PATH = '{GCS_CATEGORY_EMBEDDING_PATH}'\",\n",
        "    f\"CATG_PAD_SIZE = {MAX_NUM_CLS}\",\n",
        "    f\"INCLUDE_MASK = {INCLUDE_MASK}\",\n",
        "    f\"step_learning_rate_with_linear_warmup.init_learning_rate = {INIT_LEARNING_RATE}\",\n",
        "    f\"OUTPUT_SIZE = {OUTPUT_SIZE}\",\n",
        "    f\"get_host_evaluator.annotation_file = '{GCS_VAL_ANNOTATION_FILE}'\",\n",
        "    f\"train.pretrain_dir = '{GCS_CHECKPOINT_DIR}'\",\n",
        "    f\"DTYPE = %jnp.{DTYPE}\",\n",
        "]\n",
        "print(config_overrides)\n",
        "\n",
        "gin.parse_config_files_and_bindings(\n",
        "    [LOCAL_TRAIN_CONFIG_TEMPLATE_PATH], config_overrides, finalize_config=False\n",
        ")\n",
        "config = gin.config_str()\n",
        "\n",
        "with tf.io.gfile.GFile(GCS_TRAIN_CONFIG_PATH, \"w\") as f:\n",
        "    f.write(config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "290fd279e7f0"
      },
      "source": [
        "Create and run the training job with the model-garden JAX F-VLM training docker using the Vertex AI SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "003068435a36"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = get_job_name_with_datetime(prefix=\"jax_fvlm\")\n",
        "\n",
        "docker_args_list = [\n",
        "    f\"--output_dir={gcs_fuse_path(GCS_TRAIN_BASE_PATH)}\",\n",
        "    f\"--config_path={gcs_fuse_path(GCS_TRAIN_CONFIG_PATH)}\",\n",
        "    \"--mode=train_and_eval\",\n",
        "]\n",
        "print(docker_args_list)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9de7e802b53c"
      },
      "outputs": [],
      "source": [
        "# Click on the generated link in the output under \"View backing custom job:\" to see your run in the Cloud Console.\n",
        "# The job will run for appoximately 5 minutes in the current settings.\n",
        "job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=JOB_NAME,\n",
        "    container_uri=TRAIN_DOCKER_URI,\n",
        ")\n",
        "model = job.run(\n",
        "    args=docker_args_list,\n",
        "    base_output_dir=f\"{GCS_TRAIN_BASE_PATH}\",\n",
        "    replica_count=1,\n",
        "    machine_type=\"cloud-tpu\",\n",
        "    accelerator_type=\"TPU_V3\",\n",
        "    accelerator_count=NUM_CORES,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "798d2268f869"
      },
      "source": [
        "## Convert JAX checkpoint to TensorFlow SavedModel for inference\n",
        "\n",
        "Convert the previously fine-tuned JAX F-VLM model to a TF SavedModel for online prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f31f3e19c6ee"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = get_job_name_with_datetime(prefix=\"jax_model_conversion\")\n",
        "\n",
        "docker_args_list = [\n",
        "    f\"--input_dir={gcs_fuse_path(GCS_TRAIN_BASE_PATH)}\",\n",
        "    f\"--output_dir={gcs_fuse_path(GCS_SAVED_MODEL_DIR)}\",\n",
        "    f\"--max_num_classes={MAX_NUM_CLS}\",\n",
        "    f\"--model_name={MODEL}\",\n",
        "    f\"--include_mask={INCLUDE_MASK}\",\n",
        "]\n",
        "print(docker_args_list)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d0a4457779b1"
      },
      "outputs": [],
      "source": [
        "# Create and run the model conversion job.\n",
        "# Click on the generated link in the output under \"View backing custom job:\" to see your run in the Cloud Console.\n",
        "container_uri = MODEL_CONVERSION_DOCKER_URI\n",
        "job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=JOB_NAME,\n",
        "    container_uri=container_uri,\n",
        ")\n",
        "model_conversion_workdir = os.path.join(BUCKET_URI, JOB_NAME)\n",
        "model = job.run(\n",
        "    args=docker_args_list,\n",
        "    base_output_dir=f\"{model_conversion_workdir}\",\n",
        "    replica_count=1,\n",
        "    machine_type=\"n1-highmem-8\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ayNrua2txk0B"
      },
      "source": [
        "## Convert F-VLM SavedModel to support smaller input size\n",
        "\n",
        "The F-VLM SavedModel takes image embeddings and text embeddings as input. But you can not send these inputs directly for Vertex AI Online Prediction because there is a limit of 1.5 MB on the prediction request size. So you will first convert the SavedModel format to take jpeg bytes and text-embeddings bytes as an input instead. This modified input format will meet the 1.5 MB limit requirement.\n",
        "\n",
        "If you did not train a new model, uncomment and run the cell below to use the provided saved model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "12e3d9c0221a"
      },
      "outputs": [],
      "source": [
        "# # Path to the downloaded pre-trained saved model.\n",
        "# GCS_SAVED_MODEL_DIR = f'{F_VLM_FOLDER}/checkpoints/{MODEL.replace(\"resnet_\",\"r\")}'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ef-svu2Ix1OA"
      },
      "outputs": [],
      "source": [
        "def preprocess_jpeg_byte_string(tensor_byte_string):\n",
        "    \"\"\"Converts jpeg bytes to image embeddings as an input for the original F-VLM SavedModel.\"\"\"\n",
        "    decoded_image_tensor = tf.io.decode_jpeg(tensor_byte_string, channels=3)\n",
        "    parser_fn = inputs.get_maskrcnn_parser()\n",
        "    parser_output = parser_fn({\"image\": decoded_image_tensor})\n",
        "    image_embeddings_tensor = parser_output[\"images\"]\n",
        "    return image_embeddings_tensor\n",
        "\n",
        "\n",
        "def preprocess_text_embeddings_byte_string(tensor_byte_string):\n",
        "    \"\"\"Converts text-embeddings bytes to text-embeddings as an input for the original F-VLM SavedModel.\"\"\"\n",
        "    return tf.io.parse_tensor(tensor_byte_string, tf.float32)\n",
        "\n",
        "\n",
        "def get_serve_fn(model):\n",
        "    \"\"\"Creates a serving function for the modified SavedModel which takes jpeg bytes and text-embeddings bytes as an input.\"\"\"\n",
        "\n",
        "    @tf.function(\n",
        "        input_signature=[\n",
        "            tf.TensorSpec([None], tf.string),\n",
        "            tf.TensorSpec([None], tf.string),\n",
        "        ]\n",
        "    )\n",
        "    def serve_fn(image_jpeg_bytes_inputs, text_embeddings_bytes_inputs):\n",
        "        image_embeddings_tensor = tf.map_fn(\n",
        "            preprocess_jpeg_byte_string, image_jpeg_bytes_inputs, dtype=tf.bfloat16\n",
        "        )\n",
        "        text_embeddings_tensor = tf.map_fn(\n",
        "            preprocess_text_embeddings_byte_string,\n",
        "            text_embeddings_bytes_inputs,\n",
        "            dtype=tf.float32,\n",
        "        )\n",
        "        return model({\"image\": image_embeddings_tensor, \"text\": text_embeddings_tensor})\n",
        "\n",
        "    return serve_fn\n",
        "\n",
        "\n",
        "! rm -rf {LOCAL_CONVERTED_SAVED_MODEL_DIR}\n",
        "model = tf.saved_model.load(GCS_SAVED_MODEL_DIR)\n",
        "signatures = {\n",
        "    \"serving_default\": get_serve_fn(model=model).get_concrete_function(\n",
        "        tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec([None], tf.string)\n",
        "    )\n",
        "}\n",
        "tf.saved_model.save(model, LOCAL_CONVERTED_SAVED_MODEL_DIR, signatures=signatures)\n",
        "print(\"Saved the converted SavedModel to directory: \", LOCAL_CONVERTED_SAVED_MODEL_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1cJQEETi1jsg"
      },
      "source": [
        "Copy the local converted TF SavedModel to Cloud Storage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6hlTWKxh11dF"
      },
      "outputs": [],
      "source": [
        "! gsutil -m rm -R -f {GCS_CONVERTED_SAVED_MODEL_DIR}\n",
        "! gsutil -m cp -R {LOCAL_CONVERTED_SAVED_MODEL_DIR} {GCS_CONVERTED_SAVED_MODEL_DIR}\n",
        "! gsutil ls {GCS_CONVERTED_SAVED_MODEL_DIR}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iILhhP3TfO8B"
      },
      "source": [
        "## Run online prediction\n",
        "Run online prediction with the converted TF SavedModel."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ExIyCnKf3a94"
      },
      "source": [
        "Upload TF SavedModel and deploy it to an endpoint for prediction. This step can take up to 15 minutes to finish."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t0xYDT0BxP0W"
      },
      "outputs": [],
      "source": [
        "serving_env = {\n",
        "    \"MODEL_ID\": \"F-VLM-JAX-\",\n",
        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
        "}\n",
        "\n",
        "jax_fvlm_model = aiplatform.Model.upload(\n",
        "    display_name=\"jax_fvlm\",\n",
        "    artifact_uri=GCS_CONVERTED_SAVED_MODEL_DIR,\n",
        "    serving_container_image_uri=OPTIMIZED_TF_RUNTIME_IMAGE_URI,\n",
        "    serving_container_args=[],\n",
        "    location=REGION,\n",
        "    serving_container_environment_variables=serving_env,\n",
        ")\n",
        "\n",
        "jax_fvlm_endpoint = jax_fvlm_model.deploy(\n",
        "    deployed_model_display_name=\"jax_fvlm_deployed\",\n",
        "    traffic_split={\"0\": 100},\n",
        "    machine_type=\"n1-highmem-16\",\n",
        "    min_replica_count=1,\n",
        "    max_replica_count=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w99wNhz_3ruV"
      },
      "source": [
        "Prepare input prediction image.\n",
        "\n",
        "Note: You can modify the input image as required."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0Itg0k1s30t3"
      },
      "outputs": [],
      "source": [
        "# Local path to the prediction image.\n",
        "DEMO_IMAGE_PATH = \"./prediction_image.jpg\"\n",
        "# Download the prediction image.\n",
        "! wget -O {DEMO_IMAGE_PATH} https://cdn.pixabay.com/photo/2018/02/06/12/37/fruit-3134631_1280.jpg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B1Q7AbmJ4QxZ"
      },
      "source": [
        "Define your own categories in the cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "50109d914484"
      },
      "outputs": [],
      "source": [
        "CATEGORIES = [\n",
        "    \"kiwi\",\n",
        "    \"orange\",\n",
        "    \"lemon\",\n",
        "    \"blackberry\",\n",
        "    \"pine cone\",\n",
        "    \"red orange\",\n",
        "    \"table\",\n",
        "    \"spoon\",\n",
        "    \"pine needles\",\n",
        "    \"seed\",\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ad9b7e69dc9"
      },
      "source": [
        "Prepare jpeg bytes and text-embeddings bytes inputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qxj4Xv_DhHXj"
      },
      "outputs": [],
      "source": [
        "image_jpeg_bytes_inputs = get_jpeg_bytes(\n",
        "    local_image_path=DEMO_IMAGE_PATH, new_width=1024\n",
        ")\n",
        "\n",
        "text_embeddings = generate_text_embeddings(categories=CATEGORIES)\n",
        "text_embeddings_bytes_inputs = convert_numpy_array_to_byte_string_via_tf_tensor(\n",
        "    text_embeddings\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ys88lEkK4XDp"
      },
      "source": [
        "Use base-64 encoding followed by UTF-8 decoding to package the bytes inputs and then send them to the endpoint for prediction. The Vertex AI Prediction service will automatically convert these input strings back to bytes based on the `b64` keyword.\n",
        "\n",
        "**Note: The first prediction can take up to 2 minutes due to one time JIT compilation of the model. This may cause a timeout error below. If you get a timeout error, then wait for 2 minutes and run the prediction again. You will not get the timeout error after that.**\n",
        "The subsequent predictions take 4 seconds to finish."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mj4sqTAG4sU5"
      },
      "outputs": [],
      "source": [
        "instances_list = [\n",
        "    {\n",
        "        \"image_jpeg_bytes_inputs\": {\n",
        "            \"b64\": base64.b64encode(image_jpeg_bytes_inputs).decode(\"utf-8\")\n",
        "        },\n",
        "        \"text_embeddings_bytes_inputs\": {\n",
        "            \"b64\": base64.b64encode(text_embeddings_bytes_inputs).decode(\"utf-8\")\n",
        "        },\n",
        "    }\n",
        "]\n",
        "instances = [json_format.ParseDict(s, Value()) for s in instances_list]\n",
        "prediction_output = jax_fvlm_endpoint.predict(instances=instances).predictions[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3qC-MrN6SEs"
      },
      "source": [
        "Generate output image with predicted bounding boxes, labels, and probabilities. The output image will be saved to `OUTPUT_IMAGE_PATH`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wnPY2MFN6fL6"
      },
      "outputs": [],
      "source": [
        "generate_prediction_output_image(\n",
        "    input_image_path=DEMO_IMAGE_PATH,\n",
        "    prediction_output=prediction_output,\n",
        "    output_image_path=OUTPUT_IMAGE_PATH,\n",
        "    categories=CATEGORIES,\n",
        ")\n",
        "\n",
        "img = Image.open(OUTPUT_IMAGE_PATH)\n",
        "plt.imshow(img)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "# Delete endpoint resource.\n",
        "jax_fvlm_endpoint.delete(force=True)\n",
        "\n",
        "# Delete model resource.\n",
        "jax_fvlm_model.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created.\n",
        "delete_bucket = False\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_jax_fvlm.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}