notebooks/community/model_garden/model_garden_keras_yolov8.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 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": "TirJ-SGQseby"
      },
      "source": [
        "# Vertex AI Model Garden - Keras YOLOv8 (Finetuning)\n",
        "\n",
        "<table><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_keras_yolov8.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_keras_yolov8.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use [Keras YOLOv8](https://keras.io/api/keras_cv/models/tasks/yolo_v8_detector/) in Vertex AI Model Garden.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Run local inferences for pretrained or customized models\n",
        "\n",
        "- Deploy pretrained or customized models in Google Cloud Vertex AI\n",
        "\n",
        "- Finetune models in Google Cloud Vertex AI\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), [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.\n",
        "\n",
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the Salads category of the [OpenImages dataset](https://www.tensorflow.org/datasets/catalog/open_images_v4) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). This dataset does not require any feature engineering. The version of the dataset you will use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the bounding box locations and corresponding type of salad items in an image from a class of five items: Salad, Seafood, Tomato, Baked Goods, or Cheese."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHRlxq7tAHiv"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ISVTJUmFPoJu"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
        "\n",
        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
        "# @markdown | ----------- | ----------- | ----------- |\n",
        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
        "\n",
        "# Configs for all notebooks.\n",
        "! pip3 install --quiet keras-cv==0.9.0\n",
        "! pip3 install --quiet keras-core==0.1.0\n",
        "\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import base64\n",
        "import datetime\n",
        "import importlib\n",
        "import io\n",
        "import os\n",
        "import tempfile\n",
        "import uuid\n",
        "from typing import Dict, List, Union\n",
        "\n",
        "import keras_cv\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import yaml\n",
        "from google.cloud import aiplatform\n",
        "from google.protobuf import json_format\n",
        "from google.protobuf.struct_pb2 import Value\n",
        "from keras_cv import visualization\n",
        "from PIL import Image\n",
        "\n",
        "common_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\n",
        "\n",
        "models, endpoints = {}, {}\n",
        "\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "if not REGION:\n",
        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
        "        raise ValueError(\n",
        "            \"REGION must be set. See\"\n",
        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
        "            \" available cloud locations.\"\n",
        "        )\n",
        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"keras_yolov8\")\n",
        "\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
        "\n",
        "TRAIN_CONTAINER_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-yolov8-train\"\n",
        ")\n",
        "SERVING_CONTAINER_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest\"\n",
        ")\n",
        "\n",
        "SERVING_CONTAINER_ARGS = [\"--allow_precompilation\", \"--allow_compression\"]\n",
        "RESOLUTION = 512\n",
        "\n",
        "\n",
        "def load_img(path):\n",
        "    \"\"\"Reads image from path and return PIL.Image instance.\"\"\"\n",
        "    img = tf.io.read_file(path)\n",
        "    img = tf.image.decode_jpeg(img, channels=3)\n",
        "    return Image.fromarray(np.uint8(img)).convert(\"RGB\")\n",
        "\n",
        "\n",
        "def decode_image(image_str_tensor: tf.string) -> tf.float32:\n",
        "    \"\"\"Converts and resizes image bytes to image tensor.\"\"\"\n",
        "    image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)\n",
        "    image = tf.image.resize(image, (RESOLUTION, RESOLUTION))\n",
        "    return image\n",
        "\n",
        "\n",
        "def get_label_map(label_map_yaml_filepath):\n",
        "    \"\"\"Returns class id to label mapping given a filepath to the label map.\"\"\"\n",
        "\n",
        "    temp_dir = tempfile.TemporaryDirectory()\n",
        "    label_map_yaml_filename = os.path.basename(label_map_yaml_filepath)\n",
        "    local_metrics_path = os.path.join(temp_dir.name, label_map_yaml_filename)\n",
        "\n",
        "    ! gsutil cp $label_map_yaml_filepath $local_metrics_path\n",
        "    with open(local_metrics_path, \"r\") as input_file:\n",
        "        label_map = yaml.safe_load(input_file.read())[\"label_map\"]\n",
        "    temp_dir.cleanup()\n",
        "    return label_map\n",
        "\n",
        "\n",
        "def get_prediction_instances(test_filepath, new_width=-1):\n",
        "    \"\"\"Generate instance from image path to pass to Vertex AI Endpoint for prediction.\"\"\"\n",
        "    if new_width <= 0:\n",
        "        with tf.io.gfile.GFile(test_filepath, \"rb\") as input_file:\n",
        "            encoded_string = base64.b64encode(input_file.read()).decode(\"utf-8\")\n",
        "    else:\n",
        "        img = load_img(test_filepath)\n",
        "        width, height = img.size\n",
        "        print(\"original input image size: \", width, \" , \", height)\n",
        "        new_height = int(height * new_width / width)\n",
        "        new_img = img.resize((new_width, new_height))\n",
        "        print(\"resized input image size: \", new_width, \" , \", new_height)\n",
        "        buffered = io.BytesIO()\n",
        "        new_img.save(buffered, format=\"JPEG\")\n",
        "        encoded_string = base64.b64encode(buffered.getvalue()).decode(\"utf-8\")\n",
        "\n",
        "    instances = [\n",
        "        {\n",
        "            \"encoded_image\": {\"b64\": encoded_string},\n",
        "        }\n",
        "    ]\n",
        "    return instances\n",
        "\n",
        "\n",
        "def predict_custom_trained_model(\n",
        "    project: str,\n",
        "    endpoint_id: str,\n",
        "    instances: Union[Dict, List[Dict]],\n",
        "    location: str = \"us-central1\",\n",
        "):\n",
        "    # The AI Platform services require regional API endpoints.\n",
        "    client_options = {\"api_endpoint\": f\"{location}-aiplatform.googleapis.com\"}\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ReZHmKuBUlC2"
      },
      "outputs": [],
      "source": [
        "# @title Run local inferences with pretrained model\n",
        "\n",
        "# @markdown This section shows how to run inferences locally with YOLOv8-M pretrained on PascalVOC 2012 object detection task, which consists of 20 classes.\n",
        "\n",
        "test_filepath = \"\"  # @param {type:\"string\"}\n",
        "img_bytes = tf.io.read_file(test_filepath)\n",
        "image = tf.expand_dims(decode_image(img_bytes), axis=0)\n",
        "\n",
        "# Load model pretrained on PascalVOC 2012.\n",
        "model_name = \"yolo_v8_m_pascalvoc\"\n",
        "model = keras_cv.models.YOLOV8Detector.from_preset(\n",
        "    model_name,\n",
        "    bounding_box_format=\"xywh\",\n",
        ")\n",
        "\n",
        "decoded = model.predict(image)\n",
        "\n",
        "# Classes in PascalVOC 2012 dataset.\n",
        "class_ids = [\n",
        "    \"Aeroplane\",\n",
        "    \"Bicycle\",\n",
        "    \"Bird\",\n",
        "    \"Boat\",\n",
        "    \"Bottle\",\n",
        "    \"Bus\",\n",
        "    \"Car\",\n",
        "    \"Cat\",\n",
        "    \"Chair\",\n",
        "    \"Cow\",\n",
        "    \"Dining Table\",\n",
        "    \"Dog\",\n",
        "    \"Horse\",\n",
        "    \"Motorbike\",\n",
        "    \"Person\",\n",
        "    \"Potted Plant\",\n",
        "    \"Sheep\",\n",
        "    \"Sofa\",\n",
        "    \"Train\",\n",
        "    \"Tvmonitor\",\n",
        "    \"Total\",\n",
        "]\n",
        "class_mapping = dict(zip(range(len(class_ids)), class_ids))\n",
        "\n",
        "# Visualize the results.\n",
        "visualization.plot_bounding_box_gallery(\n",
        "    image,\n",
        "    value_range=(0, 255),\n",
        "    rows=1,\n",
        "    cols=1,\n",
        "    y_pred=decoded,\n",
        "    scale=5,\n",
        "    font_scale=0.7,\n",
        "    bounding_box_format=\"xywh\",\n",
        "    class_mapping=class_mapping,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lvUe5Y7B5uRs"
      },
      "source": [
        "## Finetune with Vertex AI Custom Training Jobs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "d_l_rCoqjqKZ"
      },
      "outputs": [],
      "source": [
        "# @title Finetune\n",
        "\n",
        "# @markdown This section shows how to finetune the Keras YOLOv8 model and deploy to Vertex AI Endpoint resource.\n",
        "\n",
        "# @markdown `input_csv_path` : The input dataset in CSV format. For further details, kindly check [AutoML Image Object Detection](https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data).\n",
        "\n",
        "input_csv_path = \"gs://cloud-samples-data/vision/salads.csv\"  # @param {type:\"string\"}\n",
        "\n",
        "# Hyperparameters\n",
        "\n",
        "# @markdown `epochs`: Number of training epochs.\n",
        "epochs = 10  # @param{type:\"integer\"}\n",
        "# @markdown `learning_rate`: The learning rate of this training job.\n",
        "learning_rate = 0.0005  # @param{type:\"number\"}\n",
        "# @markdown `fpn_depth`: The depth of the CSP blocks in the Feature Pyramid Network. This is usually 1, 2, or 3, depending on the size of your YOLOV8Detector model. We recommend using 3 for 'yolo_v8_l_backbone' and 'yolo_v8_xl_backbone'.Defaults to 2.\n",
        "fpn_depth = 3  # @param{type:\"integer\"}\n",
        "# @markdown `confidence_threshold`: Only probabilities greater than this threshold will contribute to the final result\n",
        "confidence_threshold = 0.02  # @param{type:\"number\"}\n",
        "# @markdown `iou_threshold`: Intersection over Union (IoU) is a measure that shows how well the prediction bounding box aligns with the ground truth box.\n",
        "iou_threshold = 0.3  # @param{type:\"number\"}\n",
        "# @markdown `backbone`: The pretrained backbone. [Click here](https://keras.io/api/keras_cv/models/backbones/yolo_v8/) for the full list of available backbones.\n",
        "backbone = \"yolo_v8_xl_backbone_coco\"  # @param[\"yolo_v8_xs_backbone\", \"yolo_v8_s_backbone\", \"yolo_v8_m_backbone\", \"yolo_v8_l_backbone\", \"yolo_v8_xl_backbone\", \"yolo_v8_xs_backbone_coco\", \"yolo_v8_s_backbone_coco\", \"yolo_v8_m_backbone_coco\", \"yolo_v8_l_backbone_coco\", \"yolo_v8_xl_backbone_coco\"]\n",
        "\n",
        "MACHINE_TYPE = \"n1-highmem-16\"\n",
        "ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "ACCELERATOR_COUNT = 2\n",
        "\n",
        "train_job_name = common_util.get_job_name_with_datetime(\"train_yolov8\")\n",
        "model_dir = os.path.join(STAGING_BUCKET, train_job_name)\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": MACHINE_TYPE,\n",
        "            \"accelerator_type\": ACCELERATOR_TYPE,\n",
        "            \"accelerator_count\": ACCELERATOR_COUNT,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"disk_spec\": {\n",
        "            \"boot_disk_type\": \"pd-ssd\",\n",
        "            \"boot_disk_size_gb\": 500,\n",
        "        },\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
        "            \"command\": [],\n",
        "            \"env\": [\n",
        "                {\n",
        "                    \"name\": \"RESOLUTION\",\n",
        "                    \"value\": f\"{RESOLUTION}\",\n",
        "                },\n",
        "            ],\n",
        "            \"args\": [\n",
        "                f\"--input_csv_path={input_csv_path}\",\n",
        "                f\"--output_model_dir={model_dir}\",\n",
        "                f\"--epochs={epochs}\",\n",
        "                f\"--pretrained_backbone={backbone}\",\n",
        "                f\"--fpn_depth={fpn_depth}\",\n",
        "                f\"--learning_rate={learning_rate}\",\n",
        "                f\"--confidence_threshold={confidence_threshold}\",\n",
        "                f\"--iou_threshold={iou_threshold}\",\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=ACCELERATOR_COUNT,\n",
        "    is_for_training=True,\n",
        ")\n",
        "\n",
        "# Add labels for the finetuning job.\n",
        "labels = {\n",
        "    \"mg-source\": \"notebook\",\n",
        "    \"mg-notebook-name\": \"model_garden_keras_yolov8.ipynb\".split(\".\")[0],\n",
        "}\n",
        "\n",
        "labels[\"mg-tune\"] = \"publishers-google-models-yolov8\"\n",
        "versioned_model_id = model_name.lower().replace(\"_\", \"-\")\n",
        "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
        "\n",
        "train_job = aiplatform.CustomJob(\n",
        "    display_name=train_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        "    labels=labels,\n",
        ")\n",
        "\n",
        "train_job.run()\n",
        "\n",
        "print(\"The trained model is saved in: \", model_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gnqb52XB6RIw"
      },
      "source": [
        "## Deploy and Predict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "K6rUSSKmYZJ6"
      },
      "outputs": [],
      "source": [
        "# @title Upload model\n",
        "\n",
        "upload_job_name = common_util.get_job_name_with_datetime(\"upload_yolov8\")\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=ACCELERATOR_COUNT,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
        "serving_env = {\n",
        "    \"MODEL_ID\": \"keras-yolov8\",\n",
        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
        "}\n",
        "\n",
        "models[\"yolov8\"] = aiplatform.Model.upload(\n",
        "    display_name=upload_job_name,\n",
        "    artifact_uri=model_dir,\n",
        "    serving_container_image_uri=SERVING_CONTAINER_URI,\n",
        "    serving_container_args=SERVING_CONTAINER_ARGS,\n",
        "    serving_container_environment_variables=serving_env,\n",
        "    model_garden_source_model_name=\"publishers/google/models/keras-yolov8\",\n",
        ")\n",
        "\n",
        "print(\"The model name is: \", upload_job_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "GRKrkRu7Daxw"
      },
      "outputs": [],
      "source": [
        "# @title Deploy model\n",
        "\n",
        "deploy_model_name = common_util.get_job_name_with_datetime(\"deploy_yolov8\")\n",
        "\n",
        "endpoints[\"yolov8\"] = models[\"yolov8\"].deploy(\n",
        "    deployed_model_display_name=deploy_model_name,\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    traffic_split={\"0\": 100},\n",
        "    accelerator_type=\"NVIDIA_TESLA_T4\",\n",
        "    accelerator_count=1,\n",
        "    min_replica_count=1,\n",
        "    max_replica_count=1,\n",
        "    system_labels={\"NOTEBOOK_NAME\": \"model_garden_keras_yolov8.ipynb\"},\n",
        ")\n",
        "\n",
        "\n",
        "endpoint_id = endpoints[\"yolov8\"].name\n",
        "print(\"The endpoint id is: \", endpoint_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VDznWEMmbwj4"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "test_filepath = \"gs://cloud-ml-data/img/openimage/1302/4677521502_6f2767039c_o.jpg\"  # @param {type:\"string\"}\n",
        "image_bytes = tf.io.read_file(test_filepath)\n",
        "image_resized = tf.expand_dims(decode_image(image_bytes), axis=0)\n",
        "\n",
        "instances = get_prediction_instances(test_filepath, new_width=640)\n",
        "\n",
        "predictions, _ = predict_custom_trained_model(\n",
        "    project=PROJECT_ID, location=REGION, endpoint_id=endpoint_id, instances=instances\n",
        ")\n",
        "\n",
        "predictions_dict = {\n",
        "    \"boxes\": tf.expand_dims(predictions[0][\"boxes\"], axis=0),\n",
        "    \"classes\": tf.expand_dims(predictions[0][\"classes\"], axis=0),\n",
        "    \"confidence\": tf.expand_dims(predictions[0][\"confidence\"], axis=0),\n",
        "    \"num_detections\": predictions[0][\"num_detections\"],\n",
        "}\n",
        "\n",
        "label_map = get_label_map(os.path.join(model_dir, \"label_map.yaml\"))\n",
        "\n",
        "visualization.plot_bounding_box_gallery(\n",
        "    image_resized,\n",
        "    value_range=(0, 255),\n",
        "    rows=1,\n",
        "    cols=1,\n",
        "    y_pred=predictions_dict,\n",
        "    scale=5,\n",
        "    font_scale=0.7,\n",
        "    bounding_box_format=\"xywh\",\n",
        "    class_mapping=label_map,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tWsvA3Xb6ZEm"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "# @title Delete the models and endpoints\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "for endpoint in endpoints.values():\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "for model in models.values():\n",
        "    model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1dijQDiZWegt"
      },
      "source": [
        "## References\n",
        "\n",
        "- [Efficient Object Detection with YOLOV8 and KerasCV](https://keras.io/examples/vision/yolov8/)\n",
        "- [Keras YOLOv8 API Documentation](https://keras.io/api/keras_cv/models/tasks/yolo_v8_detector/)\n",
        "- [Keras YOLOv8 Backbones](https://keras.io/api/keras_cv/models/backbones/yolo_v8/)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_keras_yolov8.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}