notebooks/community/model_garden/model_garden_pytorch_detectron2.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 - Detectron2\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_detectron2.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_pytorch_detectron2.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_detectron2.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": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates finetuning Detectron2 based [Faster R-CNN](https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md#faster-r-cnn) and\n",
        "[RetinaNet](https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md#retinanet)\n",
        "for image detection task and [Mask R-CNN](https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md#coco-instance-segmentation-baselines-with-mask-r-cnn) for segmentation task and deploying them on Vertex AI for online prediction. This notebook assumes that the input training data is in [COCO format](https://opencv.org/introduction-to-the-coco-dataset/). If you do not have your own dataset, this notebook also shows how to download and prepare the Balloon dataset for training.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Finetune a Detectron2 based Faster R-CNN, RetinaNet, or Mask R-CNN 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 image object detection and segmentation.\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": [
        "!pip3 install --upgrade google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b60a4d7100bf"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth as google_auth\n",
        "\n",
        "google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fb671e75ca7b"
      },
      "source": [
        "### Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dc8ee367fb42"
      },
      "outputs": [],
      "source": [
        "# Install gdown for downloading example training images.\n",
        "!pip install gdown\n",
        "# Install gsutil for downloading/uploading data from/to Cloud Storage buckets.\n",
        "!pip install gsutil\n",
        "# Install libraries for COCO format conversion of datasets.\n",
        "!pip install pycocotools==2.0.6\n",
        "!pip install opencv-python==4.7.0.72"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5244aac3d929"
      },
      "source": [
        "Restart the notebook kernel after installs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "567212ff53a6"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb7adab99e41"
      },
      "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.\n",
        "\n",
        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User`, `Storage Object Admin`, and `GCS Storage Bucket Owner` roles for deploying fine tuned model to Vertex AI endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6c460088b873"
      },
      "source": [
        "Fill following variables for experiments environment:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "855d6b96f291"
      },
      "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. For example 'gs://my_bucket'.\n",
        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The service account for deploying fine tuned model.\n",
        "# The service account looks like:\n",
        "# '<account_name>@<project>.iam.gserviceaccount.com'\n",
        "# Follow step 5 above to create this account.\n",
        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2cc825514deb"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b42bd4fa2b2d"
      },
      "outputs": [],
      "source": [
        "# The pre-built training docker image. It contains training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-detectron2-train\"\n",
        "# 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-detectron2-serve\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c250872074f"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XcYUGwr-AJGY"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "from typing import Dict, List, Union\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "from google.protobuf import json_format\n",
        "from google.protobuf.struct_pb2 import Value\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",
        "# Training\n",
        "def upload_model(\n",
        "    project: str,\n",
        "    location: str,\n",
        "    display_name: str,\n",
        "    serving_container_image_uri: str,\n",
        "    model_pth_file: str,\n",
        "    model_cfg_yaml_file: str,\n",
        "    test_threshold: float = 0.5,\n",
        "):\n",
        "\n",
        "    aiplatform.init(project=project, location=location)\n",
        "\n",
        "    serving_env = {\n",
        "        \"MODEL_PTH_FILE\": model_pth_file,\n",
        "        \"CONFIG_YAML_FILE\": model_cfg_yaml_file,\n",
        "        \"TEST_THRESHOLD\": test_threshold,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=display_name,\n",
        "        serving_container_image_uri=serving_container_image_uri,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/detectron2_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "    )\n",
        "\n",
        "    model.wait()\n",
        "\n",
        "    print(model.display_name)\n",
        "    print(model.resource_name)\n",
        "    return model\n",
        "\n",
        "\n",
        "def predict_custom_trained_model(\n",
        "    project: str,\n",
        "    endpoint_id: str,\n",
        "    instances: Union[Dict, List[Dict]],\n",
        "    location: str,\n",
        "    api_endpoint: str,\n",
        "):\n",
        "    # The AI Platform services require regional API endpoints.\n",
        "    client_options = {\"api_endpoint\": api_endpoint}\n",
        "    # Initialize client that will be used to create and send requests.\n",
        "    # This client only needs to be created once, and can be reused for multiple requests.\n",
        "    client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)\n",
        "    parameters_dict = {}\n",
        "    parameters = json_format.ParseDict(parameters_dict, Value())\n",
        "    endpoint = client.endpoint_path(\n",
        "        project=project, location=location, endpoint=endpoint_id\n",
        "    )\n",
        "    response = client.predict(\n",
        "        endpoint=endpoint, instances=instances, parameters=parameters\n",
        "    )\n",
        "    return response.predictions, response.deployed_model_id\n",
        "\n",
        "\n",
        "# Prediction\n",
        "import base64\n",
        "\n",
        "\n",
        "def get_prediction_instances(local_test_filepath):\n",
        "    with open(local_test_filepath, \"rb\") as input_file:\n",
        "        encoded_string = base64.b64encode(input_file.read()).decode(\"utf-8\")\n",
        "\n",
        "    instances = [\n",
        "        {\n",
        "            \"data\": {\"b64\": encoded_string},\n",
        "        }\n",
        "    ]\n",
        "    return instances\n",
        "\n",
        "\n",
        "import numpy as np\n",
        "# Mask encoding related\n",
        "import pycocotools.mask as mask_util\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",
        "import collections\n",
        "\n",
        "# Visualization\n",
        "from PIL import Image, ImageColor, ImageDraw, ImageFont\n",
        "\n",
        "\n",
        "def load_img(local_path):\n",
        "    return Image.open(local_path).convert(\"RGB\")\n",
        "\n",
        "\n",
        "def draw_bounding_box_on_image_array(\n",
        "    image,\n",
        "    ymin,\n",
        "    xmin,\n",
        "    ymax,\n",
        "    xmax,\n",
        "    color=\"red\",\n",
        "    thickness=4,\n",
        "    display_str_list=(),\n",
        "    use_normalized_coordinates=True,\n",
        "):\n",
        "    \"\"\"Adds a bounding box to an image (numpy array).\n",
        "\n",
        "    Bounding box coordinates can be specified in either absolute (pixel) or\n",
        "    normalized coordinates by setting the use_normalized_coordinates argument.\n",
        "\n",
        "    Args:\n",
        "      image: a numpy array with shape [height, width, 3].\n",
        "      ymin: ymin of bounding box.\n",
        "      xmin: xmin of bounding box.\n",
        "      ymax: ymax of bounding box.\n",
        "      xmax: xmax of bounding box.\n",
        "      color: color to draw bounding box. Default is red.\n",
        "      thickness: line thickness. Default value is 4.\n",
        "      display_str_list: list of strings to display in box\n",
        "                        (each to be shown on its own line).\n",
        "      use_normalized_coordinates: If True (default), treat coordinates\n",
        "        ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat\n",
        "        coordinates as absolute.\n",
        "    \"\"\"\n",
        "    image_pil = Image.fromarray(np.uint8(image)).convert(\"RGB\")\n",
        "    draw_bounding_box_on_image(\n",
        "        image_pil,\n",
        "        ymin,\n",
        "        xmin,\n",
        "        ymax,\n",
        "        xmax,\n",
        "        color,\n",
        "        thickness,\n",
        "        display_str_list,\n",
        "        use_normalized_coordinates,\n",
        "    )\n",
        "    np.copyto(image, np.array(image_pil))\n",
        "\n",
        "\n",
        "def draw_bounding_box_on_image(\n",
        "    image,\n",
        "    ymin,\n",
        "    xmin,\n",
        "    ymax,\n",
        "    xmax,\n",
        "    color=\"red\",\n",
        "    thickness=4,\n",
        "    display_str_list=(),\n",
        "    use_normalized_coordinates=True,\n",
        "):\n",
        "    \"\"\"Adds a bounding box to an image.\n",
        "\n",
        "    Bounding box coordinates can be specified in either absolute (pixel) or\n",
        "    normalized coordinates by setting the use_normalized_coordinates argument.\n",
        "\n",
        "    Each string in display_str_list is displayed on a separate line above the\n",
        "    bounding box in black text on a rectangle filled with the input 'color'.\n",
        "    If the top of the bounding box extends to the edge of the image, the strings\n",
        "    are displayed below the bounding box.\n",
        "\n",
        "    Args:\n",
        "      image: a PIL.Image object.\n",
        "      ymin: ymin of bounding box.\n",
        "      xmin: xmin of bounding box.\n",
        "      ymax: ymax of bounding box.\n",
        "      xmax: xmax of bounding box.\n",
        "      color: color to draw bounding box. Default is red.\n",
        "      thickness: line thickness. Default value is 4.\n",
        "      display_str_list: list of strings to display in box\n",
        "                        (each to be shown on its own line).\n",
        "      use_normalized_coordinates: If True (default), treat coordinates\n",
        "        ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat\n",
        "        coordinates as absolute.\n",
        "    \"\"\"\n",
        "    draw = ImageDraw.Draw(image)\n",
        "    im_width, im_height = image.size\n",
        "    if use_normalized_coordinates:\n",
        "        (left, right, top, bottom) = (\n",
        "            xmin * im_width,\n",
        "            xmax * im_width,\n",
        "            ymin * im_height,\n",
        "            ymax * im_height,\n",
        "        )\n",
        "    else:\n",
        "        (left, right, top, bottom) = (xmin, xmax, ymin, ymax)\n",
        "    draw.line(\n",
        "        [(left, top), (left, bottom), (right, bottom), (right, top), (left, top)],\n",
        "        width=thickness,\n",
        "        fill=color,\n",
        "    )\n",
        "    try:\n",
        "        font = ImageFont.truetype(\"arial.ttf\", 24)\n",
        "    except OSError:\n",
        "        font = ImageFont.load_default()\n",
        "\n",
        "    # If the total height of the display strings added to the top of the bounding\n",
        "    # box exceeds the top of the image, stack the strings below the bounding box\n",
        "    # instead of above.\n",
        "    display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]\n",
        "    # Each display_str has a top and bottom margin of 0.05x.\n",
        "    total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)\n",
        "\n",
        "    if top > total_display_str_height:\n",
        "        text_bottom = top\n",
        "    else:\n",
        "        text_bottom = bottom + total_display_str_height\n",
        "    # Reverse list and print from bottom to top.\n",
        "    for display_str in display_str_list[::-1]:\n",
        "        text_width, text_height = font.getsize(display_str)\n",
        "        margin = np.ceil(0.05 * text_height)\n",
        "        draw.rectangle(\n",
        "            [\n",
        "                (left, text_bottom - text_height - 2 * margin),\n",
        "                (left + text_width, text_bottom),\n",
        "            ],\n",
        "            fill=color,\n",
        "        )\n",
        "        draw.text(\n",
        "            (left + margin, text_bottom - text_height - margin),\n",
        "            display_str,\n",
        "            fill=\"black\",\n",
        "            font=font,\n",
        "        )\n",
        "        text_bottom -= text_height - 2 * margin\n",
        "\n",
        "\n",
        "def draw_mask_on_image_array(image, mask, color=\"red\", alpha=0.4):\n",
        "    \"\"\"Draws mask on an image.\n",
        "\n",
        "    Args:\n",
        "      image: uint8 numpy array with shape (img_height, img_height, 3)\n",
        "      mask: a uint8 numpy array of shape (img_height, img_height) with\n",
        "        values between either 0 or 1.\n",
        "      color: color to draw the keypoints with. Default is red.\n",
        "      alpha: transparency value between 0 and 1. (default: 0.4)\n",
        "\n",
        "    Raises:\n",
        "      ValueError: On incorrect data type for image or masks.\n",
        "    \"\"\"\n",
        "    if image.dtype != np.uint8:\n",
        "        raise ValueError(\"`image` not of type np.uint8\")\n",
        "    if mask.dtype != np.uint8:\n",
        "        raise ValueError(\"`mask` not of type np.uint8\")\n",
        "    if np.any(np.logical_and(mask != 1, mask != 0)):\n",
        "        raise ValueError(\"`mask` elements should be in [0, 1]\")\n",
        "    if image.shape[:2] != mask.shape:\n",
        "        raise ValueError(\n",
        "            \"The image has spatial dimensions %s but the mask has \"\n",
        "            \"dimensions %s\" % (image.shape[:2], mask.shape)\n",
        "        )\n",
        "    rgb = ImageColor.getrgb(color)\n",
        "    pil_image = Image.fromarray(image)\n",
        "\n",
        "    solid_color = np.expand_dims(np.ones_like(mask), axis=2) * np.reshape(\n",
        "        list(rgb), [1, 1, 3]\n",
        "    )\n",
        "    pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert(\"RGBA\")\n",
        "    pil_mask = Image.fromarray(np.uint8(255.0 * alpha * mask)).convert(\"L\")\n",
        "    pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)\n",
        "    np.copyto(image, np.array(pil_image.convert(\"RGB\")))\n",
        "\n",
        "\n",
        "STANDARD_COLORS = [\n",
        "    \"AliceBlue\",\n",
        "    \"Chartreuse\",\n",
        "    \"Aqua\",\n",
        "    \"Aquamarine\",\n",
        "    \"Azure\",\n",
        "    \"Beige\",\n",
        "    \"Bisque\",\n",
        "    \"BlanchedAlmond\",\n",
        "    \"BlueViolet\",\n",
        "    \"BurlyWood\",\n",
        "    \"CadetBlue\",\n",
        "    \"AntiqueWhite\",\n",
        "    \"Chocolate\",\n",
        "    \"Coral\",\n",
        "    \"CornflowerBlue\",\n",
        "    \"Cornsilk\",\n",
        "    \"Crimson\",\n",
        "    \"Cyan\",\n",
        "    \"DarkCyan\",\n",
        "    \"DarkGoldenRod\",\n",
        "    \"DarkGrey\",\n",
        "    \"DarkKhaki\",\n",
        "    \"DarkOrange\",\n",
        "    \"DarkOrchid\",\n",
        "    \"DarkSalmon\",\n",
        "    \"DarkSeaGreen\",\n",
        "    \"DarkTurquoise\",\n",
        "    \"DarkViolet\",\n",
        "    \"DeepPink\",\n",
        "    \"DeepSkyBlue\",\n",
        "    \"DodgerBlue\",\n",
        "    \"FireBrick\",\n",
        "    \"FloralWhite\",\n",
        "    \"ForestGreen\",\n",
        "    \"Fuchsia\",\n",
        "    \"Gainsboro\",\n",
        "    \"GhostWhite\",\n",
        "    \"Gold\",\n",
        "    \"GoldenRod\",\n",
        "    \"Salmon\",\n",
        "    \"Tan\",\n",
        "    \"HoneyDew\",\n",
        "    \"HotPink\",\n",
        "    \"IndianRed\",\n",
        "    \"Ivory\",\n",
        "    \"Khaki\",\n",
        "    \"Lavender\",\n",
        "    \"LavenderBlush\",\n",
        "    \"LawnGreen\",\n",
        "    \"LemonChiffon\",\n",
        "    \"LightBlue\",\n",
        "    \"LightCoral\",\n",
        "    \"LightCyan\",\n",
        "    \"LightGoldenRodYellow\",\n",
        "    \"LightGray\",\n",
        "    \"LightGrey\",\n",
        "    \"LightGreen\",\n",
        "    \"LightPink\",\n",
        "    \"LightSalmon\",\n",
        "    \"LightSeaGreen\",\n",
        "    \"LightSkyBlue\",\n",
        "    \"LightSlateGray\",\n",
        "    \"LightSlateGrey\",\n",
        "    \"LightSteelBlue\",\n",
        "    \"LightYellow\",\n",
        "    \"Lime\",\n",
        "    \"LimeGreen\",\n",
        "    \"Linen\",\n",
        "    \"Magenta\",\n",
        "    \"MediumAquaMarine\",\n",
        "    \"MediumOrchid\",\n",
        "    \"MediumPurple\",\n",
        "    \"MediumSeaGreen\",\n",
        "    \"MediumSlateBlue\",\n",
        "    \"MediumSpringGreen\",\n",
        "    \"MediumTurquoise\",\n",
        "    \"MediumVioletRed\",\n",
        "    \"MintCream\",\n",
        "    \"MistyRose\",\n",
        "    \"Moccasin\",\n",
        "    \"NavajoWhite\",\n",
        "    \"OldLace\",\n",
        "    \"Olive\",\n",
        "    \"OliveDrab\",\n",
        "    \"Orange\",\n",
        "    \"OrangeRed\",\n",
        "    \"Orchid\",\n",
        "    \"PaleGoldenRod\",\n",
        "    \"PaleGreen\",\n",
        "    \"PaleTurquoise\",\n",
        "    \"PaleVioletRed\",\n",
        "    \"PapayaWhip\",\n",
        "    \"PeachPuff\",\n",
        "    \"Peru\",\n",
        "    \"Pink\",\n",
        "    \"Plum\",\n",
        "    \"PowderBlue\",\n",
        "    \"Purple\",\n",
        "    \"Red\",\n",
        "    \"RosyBrown\",\n",
        "    \"RoyalBlue\",\n",
        "    \"SaddleBrown\",\n",
        "    \"Green\",\n",
        "    \"SandyBrown\",\n",
        "    \"SeaGreen\",\n",
        "    \"SeaShell\",\n",
        "    \"Sienna\",\n",
        "    \"Silver\",\n",
        "    \"SkyBlue\",\n",
        "    \"SlateBlue\",\n",
        "    \"SlateGray\",\n",
        "    \"SlateGrey\",\n",
        "    \"Snow\",\n",
        "    \"SpringGreen\",\n",
        "    \"SteelBlue\",\n",
        "    \"GreenYellow\",\n",
        "    \"Teal\",\n",
        "    \"Thistle\",\n",
        "    \"Tomato\",\n",
        "    \"Turquoise\",\n",
        "    \"Violet\",\n",
        "    \"Wheat\",\n",
        "    \"White\",\n",
        "    \"WhiteSmoke\",\n",
        "    \"Yellow\",\n",
        "    \"YellowGreen\",\n",
        "]\n",
        "\n",
        "\n",
        "def visualize_boxes_and_labels_on_image_array(\n",
        "    image,\n",
        "    boxes,\n",
        "    classes,\n",
        "    scores,\n",
        "    category_index,\n",
        "    instance_masks=None,\n",
        "    use_normalized_coordinates=False,\n",
        "    max_boxes_to_draw=20,\n",
        "    min_score_thresh=0.5,\n",
        "    agnostic_mode=False,\n",
        "    line_thickness=4,\n",
        "    groundtruth_box_visualization_color=\"black\",\n",
        "    skip_scores=False,\n",
        "    skip_labels=False,\n",
        "):\n",
        "    \"\"\"Overlay labeled boxes on an image with formatted scores and label names.\n",
        "\n",
        "    This function groups boxes that correspond to the same location\n",
        "    and creates a display string for each detection and overlays these\n",
        "    on the image. Note that this function modifies the image in place, and returns\n",
        "    that same image.\n",
        "\n",
        "    Args:\n",
        "      image: uint8 numpy array with shape (img_height, img_width, 3)\n",
        "      boxes: a numpy array of shape [N, 4]\n",
        "      classes: a numpy array of shape [N]. Note that class indices are 1-based,\n",
        "        and match the keys in the label map.\n",
        "      scores: a numpy array of shape [N] or None.  If scores=None, then\n",
        "        this function assumes that the boxes to be plotted are groundtruth\n",
        "        boxes and plot all boxes as black with no classes or scores.\n",
        "      category_index: a dict containing category dictionaries (each holding\n",
        "        category index `id` and category name `name`) keyed by category indices.\n",
        "      instance_masks: a numpy array of shape [N, image_height, image_width] with\n",
        "        values ranging between 0 and 1, can be None.\n",
        "      instance_boundaries: a numpy array of shape [N, image_height, image_width]\n",
        "        with values ranging between 0 and 1, can be None.\n",
        "      keypoints: a numpy array of shape [N, num_keypoints, 2], can\n",
        "        be None\n",
        "      use_normalized_coordinates: whether boxes is to be interpreted as\n",
        "        normalized coordinates or not.\n",
        "      max_boxes_to_draw: maximum number of boxes to visualize.  If None, draw\n",
        "        all boxes.\n",
        "      min_score_thresh: minimum score threshold for a box to be visualized\n",
        "      agnostic_mode: boolean (default: False) controlling whether to evaluate in\n",
        "        class-agnostic mode or not.  This mode will display scores but ignore\n",
        "        classes.\n",
        "      line_thickness: integer (default: 4) controlling line width of the boxes.\n",
        "      groundtruth_box_visualization_color: box color for visualizing groundtruth\n",
        "        boxes\n",
        "      skip_scores: whether to skip score when drawing a single detection\n",
        "      skip_labels: whether to skip label when drawing a single detection\n",
        "\n",
        "    Returns:\n",
        "      uint8 numpy array with shape (img_height, img_width, 3) with overlaid boxes.\n",
        "    \"\"\"\n",
        "    # Create a display string (and color) for every box location, group any boxes\n",
        "    # that correspond to the same location.\n",
        "    box_to_display_str_map = collections.defaultdict(list)\n",
        "    box_to_color_map = collections.defaultdict(str)\n",
        "    box_to_instance_masks_map = {}\n",
        "    if not max_boxes_to_draw:\n",
        "        max_boxes_to_draw = boxes.shape[0]\n",
        "    for i in range(min(max_boxes_to_draw, boxes.shape[0])):\n",
        "        if scores is None or scores[i] > min_score_thresh:\n",
        "            box = tuple(boxes[i].tolist())\n",
        "            if instance_masks is not None:\n",
        "                box_to_instance_masks_map[box] = instance_masks[i]\n",
        "            if scores is None:\n",
        "                box_to_color_map[box] = groundtruth_box_visualization_color\n",
        "            else:\n",
        "                display_str = \"\"\n",
        "                if not skip_labels:\n",
        "                    if not agnostic_mode:\n",
        "                        if classes[i] in category_index.keys():\n",
        "                            class_name = category_index[classes[i]][\"name\"]\n",
        "                        else:\n",
        "                            class_name = \"N/A\"\n",
        "                        display_str = str(class_name)\n",
        "                if not skip_scores:\n",
        "                    if not display_str:\n",
        "                        display_str = \"{}%\".format(int(100 * scores[i]))\n",
        "                    else:\n",
        "                        display_str = \"{}: {}%\".format(\n",
        "                            display_str, int(100 * scores[i])\n",
        "                        )\n",
        "                box_to_display_str_map[box].append(display_str)\n",
        "                if agnostic_mode:\n",
        "                    box_to_color_map[box] = \"DarkOrange\"\n",
        "                else:\n",
        "                    box_to_color_map[box] = STANDARD_COLORS[\n",
        "                        classes[i] % len(STANDARD_COLORS)\n",
        "                    ]\n",
        "\n",
        "    # Draw all boxes onto image.\n",
        "    for box, color in box_to_color_map.items():\n",
        "        # Using Detectron2 style output.\n",
        "        xmin, ymin, xmax, ymax = box\n",
        "        if instance_masks is not None:\n",
        "            draw_mask_on_image_array(image, box_to_instance_masks_map[box], color=color)\n",
        "        draw_bounding_box_on_image_array(\n",
        "            image,\n",
        "            ymin,\n",
        "            xmin,\n",
        "            ymax,\n",
        "            xmax,\n",
        "            color=color,\n",
        "            thickness=line_thickness,\n",
        "            display_str_list=box_to_display_str_map[box],\n",
        "            use_normalized_coordinates=use_normalized_coordinates,\n",
        "        )\n",
        "\n",
        "    return image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RB_xY9ipr7ZU"
      },
      "source": [
        "### (Optional) Download and prepare Balloon dataset\n",
        "\n",
        "You only need this step if you do not have your own dataset and want to use the Balloon dataset as a demo. If using your own dataset, convert it into [COCO format](https://opencv.org/introduction-to-the-coco-dataset/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RcoOH0yEpJQ7"
      },
      "outputs": [],
      "source": [
        "# Download Balloon data.\n",
        "!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip\n",
        "!unzip balloon_dataset.zip > /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E5StJ42NpJQ8"
      },
      "outputs": [],
      "source": [
        "local_balloon_data_directory = \"balloon\"  # @param {type:\"string\"}\n",
        "BALLOON_DATA_GCS_PATH = os.path.join(BUCKET_URI, \"balloon_dataset\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ah0lXyvXpJQ8"
      },
      "outputs": [],
      "source": [
        "# Convert Balloon data to COCO format\n",
        "import cv2\n",
        "\n",
        "\n",
        "def save_coco_format_json(img_dir, output_coco_format_json_filename):\n",
        "    # Load original balloon data json file\n",
        "    json_file = os.path.join(img_dir, \"via_region_data.json\")\n",
        "    with open(json_file) as f:\n",
        "        imgs_anns = json.load(f)\n",
        "\n",
        "    output_coco_format_dict = {}\n",
        "    # We only have one class: balloon.\n",
        "    output_coco_format_dict[\"categories\"] = [{\"id\": 0, \"name\": \"balloon\"}]\n",
        "    output_coco_format_dict[\"images\"] = []\n",
        "    output_coco_format_dict[\"annotations\"] = []\n",
        "    annotation_idx = 0\n",
        "    for image_idx, v in enumerate(imgs_anns.values()):\n",
        "        filename = os.path.join(img_dir, v[\"filename\"])\n",
        "        height, width = cv2.imread(filename).shape[:2]\n",
        "        image_item = {\n",
        "            \"id\": image_idx,\n",
        "            \"width\": width,\n",
        "            \"height\": height,\n",
        "            \"file_name\": v[\"filename\"],\n",
        "        }\n",
        "        output_coco_format_dict[\"images\"].append(image_item)\n",
        "\n",
        "        # Process all regions in this image.\n",
        "        annos = v[\"regions\"]\n",
        "        for _, anno in annos.items():\n",
        "            assert not anno[\"region_attributes\"]\n",
        "            anno = anno[\"shape_attributes\"]\n",
        "            px = anno[\"all_points_x\"]\n",
        "            py = anno[\"all_points_y\"]\n",
        "            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]\n",
        "            poly = [p for x in poly for p in x]\n",
        "\n",
        "            annotation_idx += 1\n",
        "            annotation_item = {\n",
        "                \"id\": annotation_idx,\n",
        "                \"image_id\": image_idx,\n",
        "                # x, y, width, height\n",
        "                \"bbox\": [\n",
        "                    int(np.min(px)),\n",
        "                    int(np.min(py)),\n",
        "                    int(np.max(px) - np.min(px)),\n",
        "                    int(np.max(py) - np.min(py)),\n",
        "                ],\n",
        "                \"iscrowd\": 0,\n",
        "                # Only have one category.\n",
        "                \"category_id\": 0,\n",
        "                \"segmentation\": [poly],\n",
        "            }\n",
        "            RLEs = mask_util.frPyObjects([poly], width, height)\n",
        "            RLE = mask_util.merge(RLEs)\n",
        "            annotation_item[\"area\"] = float(mask_util.area(RLE))\n",
        "            output_coco_format_dict[\"annotations\"].append(annotation_item)\n",
        "\n",
        "    # Save output file.\n",
        "    json_file = os.path.join(img_dir, output_coco_format_json_filename)\n",
        "    with open(json_file, \"w\") as f:\n",
        "        json.dump(output_coco_format_dict, f)\n",
        "\n",
        "\n",
        "save_coco_format_json(\n",
        "    os.path.join(local_balloon_data_directory, \"train\"),\n",
        "    \"balloon_train_coco_format.json\",\n",
        ")\n",
        "save_coco_format_json(\n",
        "    os.path.join(local_balloon_data_directory, \"val\"), \"balloon_val_coco_format.json\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0OykIZen9gCC"
      },
      "outputs": [],
      "source": [
        "# Move Balloon data from local directory to Cloud Storage.\n",
        "\n",
        "import glob\n",
        "\n",
        "from google.cloud import storage\n",
        "\n",
        "\n",
        "def get_bucket_and_blob_name(filepath):\n",
        "    # The gcs path is of the form gs://<bucket-name>/<blob-name>\n",
        "    gs_suffix = filepath.split(\"gs://\", 1)[1]\n",
        "    return tuple(gs_suffix.split(\"/\", 1))\n",
        "\n",
        "\n",
        "def upload_local_dir_to_gcs(local_dir_path, gcs_dir_path):\n",
        "    \"\"\"Uploads files in a local directory to a GCS directory.\"\"\"\n",
        "    client = storage.Client()\n",
        "    bucket_name = gcs_dir_path.split(\"/\")[2]\n",
        "    bucket = client.get_bucket(bucket_name)\n",
        "    for local_file in glob.glob(local_dir_path + \"/**\"):\n",
        "        if not os.path.isfile(local_file):\n",
        "            continue\n",
        "        filename = local_file[1 + len(local_dir_path) :]\n",
        "        gcs_file_path = os.path.join(gcs_dir_path, filename)\n",
        "        _, blob_name = get_bucket_and_blob_name(gcs_file_path)\n",
        "        blob = bucket.blob(blob_name)\n",
        "        blob.upload_from_filename(local_file)\n",
        "        print(\"Copied {} to {}.\".format(local_file, gcs_file_path))\n",
        "\n",
        "\n",
        "upload_local_dir_to_gcs(\n",
        "    os.path.join(local_balloon_data_directory, \"train\"),\n",
        "    os.path.join(BALLOON_DATA_GCS_PATH, \"train\"),\n",
        ")\n",
        "upload_local_dir_to_gcs(\n",
        "    os.path.join(local_balloon_data_directory, \"val\"),\n",
        "    os.path.join(BALLOON_DATA_GCS_PATH, \"val\"),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aCpLmWPMpJQ8"
      },
      "source": [
        "## Finetune with Detectron2\n",
        "\n",
        "You will use the Vertex AI SDK to create and run the training job with the model-garden detectron2 training docker. You can choose one of the Faster R-CNN, RetinaNet, or Mask R-CNN models to finetune by uncommenting the corresponding code sections below. The training uses one V100 GPU and runs for around 3 mins once the training job begins."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aec22792ee84"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
        "JOB_NAME = \"detectron2_balloon_\" + TIMESTAMP\n",
        "\n",
        "NUM_GPU = 1\n",
        "container_uri = TRAIN_DOCKER_URI\n",
        "staging_bucket = os.path.join(BUCKET_URI, \"training/temporal\")\n",
        "\n",
        "# Dataset and output directory related parameters.\n",
        "train_dataset_name = \"balloon_train\"  # @param {type:\"string\"}\n",
        "train_coco_json_file = os.path.join(\n",
        "    BALLOON_DATA_GCS_PATH, \"train/balloon_train_coco_format.json\"\n",
        ")\n",
        "train_coco_json_file = gcs_fuse_path(train_coco_json_file)\n",
        "train_image_root = os.path.join(BALLOON_DATA_GCS_PATH, \"train\")\n",
        "train_image_root = gcs_fuse_path(train_image_root)\n",
        "val_dataset_name = \"balloon_val\"  # @param {type:\"string\"}\n",
        "val_coco_json_file = os.path.join(\n",
        "    BALLOON_DATA_GCS_PATH, \"val/balloon_val_coco_format.json\"\n",
        ")\n",
        "val_coco_json_file = gcs_fuse_path(val_coco_json_file)\n",
        "val_image_root = os.path.join(BALLOON_DATA_GCS_PATH, \"val\")\n",
        "val_image_root = gcs_fuse_path(val_image_root)\n",
        "output_dir = os.path.join(BUCKET_URI, JOB_NAME)\n",
        "\n",
        "#################################################\n",
        "# Model and dataset related parameters for Mask R-CNN.\n",
        "config_file = \"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"\n",
        "config_file = gcs_fuse_path(config_file)\n",
        "remainder_args_list = []\n",
        "remainder_args_list += [\"DATASETS.TRAIN\"] + [\n",
        "    '(\"{train_dataset_name}\",)'.format(train_dataset_name=train_dataset_name)\n",
        "]\n",
        "remainder_args_list += [\"DATASETS.TEST\"] + [\n",
        "    '(\"{val_dataset_name}\",)'.format(val_dataset_name=val_dataset_name)\n",
        "]\n",
        "remainder_args_list += [\"DATALOADER.NUM_WORKERS\"] + [\"2\"]\n",
        "remainder_args_list += [\"SOLVER.IMS_PER_BATCH\"] + [\"2\"]\n",
        "remainder_args_list += [\"SOLVER.MAX_ITER\"] + [\"300\"]\n",
        "remainder_args_list += [\"SOLVER.STEPS\"] + [\"[]\"]\n",
        "remainder_args_list += [\"MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE\"] + [\"128\"]\n",
        "remainder_args_list += [\"MODEL.ROI_HEADS.NUM_CLASSES\"] + [\"1\"]\n",
        "#################################################\n",
        "\n",
        "# #################################################\n",
        "# # Model and dataset related parameters for RetinaNet.\n",
        "# config_file='COCO-Detection/retinanet_R_50_FPN_3x.yaml'\n",
        "# config_file = gcs_fuse_path(config_file)\n",
        "# remainder_args_list = []\n",
        "# remainder_args_list += ['DATASETS.TRAIN'] + ['(\"{train_dataset_name}\",)'.format(train_dataset_name=train_dataset_name)]\n",
        "# remainder_args_list += ['DATASETS.TEST'] + ['(\"{val_dataset_name}\",)'.format(val_dataset_name=val_dataset_name)]\n",
        "# remainder_args_list += ['DATALOADER.NUM_WORKERS'] + ['2']\n",
        "# remainder_args_list += ['SOLVER.IMS_PER_BATCH'] + ['2']\n",
        "# remainder_args_list += ['SOLVER.MAX_ITER'] + ['300']\n",
        "# remainder_args_list += ['SOLVER.STEPS'] + ['[]']\n",
        "# remainder_args_list += ['MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE'] + ['128']\n",
        "# remainder_args_list += ['MODEL.ROI_HEADS.NUM_CLASSES'] + ['1']\n",
        "# remainder_args_list += ['MODEL.RETINANET.NUM_CLASSES'] + ['1']\n",
        "# #################################################\n",
        "\n",
        "# #################################################\n",
        "# # Model and dataset related parameters for Faster R-CNN.\n",
        "# config_file='COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml'\n",
        "# config_file = gcs_fuse_path(config_file)\n",
        "# remainder_args_list = []\n",
        "# remainder_args_list += ['DATASETS.TRAIN'] + ['(\"{train_dataset_name}\",)'.format(train_dataset_name=train_dataset_name)]\n",
        "# remainder_args_list += ['DATASETS.TEST'] + ['(\"{val_dataset_name}\",)'.format(val_dataset_name=val_dataset_name)]\n",
        "# remainder_args_list += ['DATALOADER.NUM_WORKERS'] + ['2']\n",
        "# remainder_args_list += ['SOLVER.IMS_PER_BATCH'] + ['2']\n",
        "# remainder_args_list += ['SOLVER.MAX_ITER'] + ['300']\n",
        "# remainder_args_list += ['SOLVER.STEPS'] + ['[]']\n",
        "# remainder_args_list += ['MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE'] + ['128']\n",
        "# remainder_args_list += ['MODEL.ROI_HEADS.NUM_CLASSES'] + ['1']\n",
        "# #################################################\n",
        "\n",
        "# Create argument list for docker.\n",
        "# NOTE: Config file flag name has hyphen instead\n",
        "# of underscore: 'config-file'.\n",
        "lr = 0.00025\n",
        "docker_args_list = [\n",
        "    \"--train_dataset_name\",\n",
        "    f\"{train_dataset_name}\",\n",
        "    \"--train_coco_json_file\",\n",
        "    f\"{train_coco_json_file}\",\n",
        "    \"--train_image_root\",\n",
        "    f\"{train_image_root}\",\n",
        "    \"--val_dataset_name\",\n",
        "    f\"{val_dataset_name}\",\n",
        "    \"--val_coco_json_file\",\n",
        "    f\"{val_coco_json_file}\",\n",
        "    \"--val_image_root\",\n",
        "    f\"{val_image_root}\",\n",
        "    \"--lr\",\n",
        "    f\"{lr}\",\n",
        "    \"--num-gpus\",\n",
        "    f\"{NUM_GPU}\",\n",
        "    \"--output_dir\",\n",
        "    f\"{gcs_fuse_path(output_dir)}\",\n",
        "    \"--config-file\",\n",
        "    f\"{config_file}\",\n",
        "]\n",
        "docker_args_list += remainder_args_list"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ELphfgj1f3Q"
      },
      "outputs": [],
      "source": [
        "# Create and run training job.\n",
        "# Click on the generated link in the output under \"View backing custom job:\" to see your run in the Cloud Console.\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=staging_bucket)\n",
        "job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=JOB_NAME,\n",
        "    container_uri=container_uri,\n",
        ")\n",
        "model = job.run(\n",
        "    args=docker_args_list,\n",
        "    base_output_dir=f\"{output_dir}\",\n",
        "    replica_count=1,\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "    accelerator_count=NUM_GPU,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g0BGaofgsMsy"
      },
      "source": [
        "## Upload and deploy Models\n",
        "\n",
        "This section uploads the model to Model Registry and deploys it on an Endpoint resource. The model deployment step will take ~15 minutes to complete. You need to set the model path below from the training output Cloud Storage directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3R_3DgCRD3U-"
      },
      "outputs": [],
      "source": [
        "# Upload models to model registry\n",
        "from datetime import datetime\n",
        "\n",
        "PRETRAINED_MODEL_PTH_FILE = os.path.join(output_dir, \"model_final.pth\")\n",
        "PRETRAINED_MODEL_CFG_YAML_FILE = os.path.join(output_dir, \"config.yaml\")\n",
        "TEST_THRESHOLD = 0.7\n",
        "PREDICTION_CONTAINER_URI = SERVE_DOCKER_URI\n",
        "PREDICTION_DISPLAY_NAME = \"upload_detectron2_\" + datetime.now().strftime(\n",
        "    \"%Y%m%d_%H%M%S\"\n",
        ")\n",
        "\n",
        "model = upload_model(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    display_name=PREDICTION_DISPLAY_NAME,\n",
        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
        "    model_pth_file=PRETRAINED_MODEL_PTH_FILE,\n",
        "    model_cfg_yaml_file=PRETRAINED_MODEL_CFG_YAML_FILE,\n",
        "    test_threshold=TEST_THRESHOLD,\n",
        ")\n",
        "print(\"The uploaded model name is: \", PREDICTION_DISPLAY_NAME)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYuQowyZEtxK"
      },
      "outputs": [],
      "source": [
        "# Deploy uploaded models\n",
        "from datetime import datetime\n",
        "\n",
        "DEPLOYED_NAME = \"deploy_iod_\" + datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
        "MACHINE_TYPE = \"n1-highmem-16\"  # @param {type:\"string\"}\n",
        "\n",
        "TRAFFIC_SPLIT = {\"0\": 100}\n",
        "\n",
        "MIN_NODES = 1\n",
        "MAX_NODES = 1\n",
        "\n",
        "endpoint = model.deploy(\n",
        "    deployed_model_display_name=DEPLOYED_NAME,\n",
        "    traffic_split=TRAFFIC_SPLIT,\n",
        "    machine_type=MACHINE_TYPE,\n",
        "    min_replica_count=MIN_NODES,\n",
        "    max_replica_count=MAX_NODES,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        ")\n",
        "\n",
        "print(\"endpoint id is: \", endpoint.name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vbIW9me1F2RY"
      },
      "outputs": [],
      "source": [
        "# Run predictions\n",
        "# Fill the \"endpoint_id\" from previous step below.\n",
        "# For example 'endpoint_id = \"8211918096324100096\"'.\n",
        "\n",
        "endpoint_id = \"\"  # @param {type:\"string\"}\n",
        "local_test_filepath = os.path.join(\n",
        "    local_balloon_data_directory, \"val/410488422_5f8991f26e_b.jpg\"\n",
        ")\n",
        "instances = get_prediction_instances(local_test_filepath)\n",
        "api_endpoint = REGION + \"-aiplatform.googleapis.com\"\n",
        "\n",
        "predictions, deployed_model_id = predict_custom_trained_model(\n",
        "    project=PROJECT_ID,\n",
        "    endpoint_id=endpoint_id,\n",
        "    instances=instances,\n",
        "    location=REGION,\n",
        "    api_endpoint=api_endpoint,\n",
        ")\n",
        "\n",
        "print(\"The deployed model id: \", deployed_model_id)\n",
        "print(\"Predict the test image: \", local_test_filepath)\n",
        "prediction = json.loads(predictions[0])\n",
        "print(prediction)\n",
        "\n",
        "# Draw boxes and masks.\n",
        "img = load_img(local_test_filepath)\n",
        "\n",
        "boxes = prediction[\"boxes\"]\n",
        "classes = prediction[\"classes\"]\n",
        "scores = prediction[\"scores\"]\n",
        "if prediction[\"masks_rle\"]:\n",
        "    masks_numpy = decode_rle_masks(prediction[\"masks_rle\"])\n",
        "else:\n",
        "    masks_numpy = None\n",
        "img.save(\"./sample.jpg\")\n",
        "output_image_array = visualize_boxes_and_labels_on_image_array(\n",
        "    image=np.array(img),\n",
        "    boxes=np.array(boxes),\n",
        "    classes=np.array(classes),\n",
        "    scores=np.array(scores),\n",
        "    category_index={0: {\"name\": \"balloon\"}},\n",
        "    instance_masks=masks_numpy,\n",
        ")\n",
        "output_image = Image.fromarray(np.uint8(output_image_array))\n",
        "output_image.save(\"./sample_preds.jpg\")\n",
        "print('Prediction image saved to \"./sample_preds.jpg\" ')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aP9ZYHaeYCaf"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aP9ZYHaeYCas"
      },
      "outputs": [],
      "source": [
        "# Cleans up endpoint from previous step.\n",
        "# Or you can overwrite \"endpoint_id\" for a different one.\n",
        "\n",
        "aip_endpoint_name = f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
        "endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "# Undeploy model and delete endpoint\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete the model resource\n",
        "model.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_detectron2.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}