guides/linked/Earth_Engine_TensorFlow_Vertex_AI.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "fSIfBsgi8dNK"
      },
      "source": [
        "#@title Copyright 2023 Google LLC. { display-mode: \"form\" }\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."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aV1xZ1CPi3Nw"
      },
      "source": [
        "\u003ctable class=\"ee-notebook-buttons\" align=\"left\"\u003e\u003ctd\u003e\n",
        "\u003ca target=\"_blank\"  href=\"http://colab.research.google.com/github/google/earthengine-community/blob/master/guides/linked/Earth_Engine_TensorFlow_Vertex_AI.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003e Run in Google Colab\u003c/a\u003e\n",
        "\u003c/td\u003e\u003ctd\u003e\n",
        "\u003ca target=\"_blank\"  href=\"https://github.com/google/earthengine-community/blob/master/guides/linked/Earth_Engine_TensorFlow_Vertex_AI.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003e View source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AC8adBmw-5m3"
      },
      "source": [
        "# Connecting Earth Engine to a Vertex AI hosted model\n",
        "\n",
        "This notebook demonstrates training a per-pixel neural network in TensorFlow, hosting the model on Vertex AI and using it in Earth Engine for interactive prediction from an `ee.Model.fromVertexAIPredictor`.\n",
        "\n",
        "**Running this demo may incur charges to your Google Cloud Account!**"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import auth\n",
        "from pprint import pprint\n",
        "\n",
        "import ee\n",
        "import folium\n",
        "import google\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras"
      ],
      "metadata": {
        "id": "E-nnPe9x8BBn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Authentication and Initialization"
      ],
      "metadata": {
        "id": "uI73xMI99aO8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "auth.authenticate_user()"
      ],
      "metadata": {
        "id": "TLmI05-wT_GD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# REPLACE WITH YOUR CLOUD PROJECT!\n",
        "PROJECT = 'your-project'"
      ],
      "metadata": {
        "id": "xESt17S3M3PF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "credentials, _ = google.auth.default()\n",
        "ee.Initialize(credentials, project=PROJECT, opt_url='https://earthengine-highvolume.googleapis.com')"
      ],
      "metadata": {
        "id": "c5bEkwQHUDPS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DrXLkJC2QJdP"
      },
      "source": [
        "# Define variables\n",
        "\n",
        "The training data are land cover labels with a single vector of Landsat 8 pixel values (`BANDS`) as predictors.  See [this example notebook](http://colab.research.google.com/github/google/earthengine-api/blob/master/python/examples/ipynb/Earth_Engine_TensorFlow_DNN_from_scratch.ipynb) for details on how to generate these training data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GHTOc5YLQZ5B"
      },
      "source": [
        "# Output bucket for trained models. REPLACE WITH YOUR WRITABLE BUCKET!\n",
        "OUTPUT_BUCKET = 'your-bucket'\n",
        "\n",
        "REGION = 'us-central1'\n",
        "\n",
        "# Cloud Storage bucket with training and testing datasets.\n",
        "DATA_BUCKET = 'ee-docs-demos'\n",
        "\n",
        "MODEL_NAME = 'vertex_dnn_demo'\n",
        "MODEL_DIR = 'gs://' + OUTPUT_BUCKET + '/' + MODEL_NAME\n",
        "CONTAINER_IMAGE = 'us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest'\n",
        "\n",
        "ENDPOINT_NAME = 'dnn_demo_endpoint'\n",
        "\n",
        "# Training and testing dataset file names in the Cloud Storage bucket.\n",
        "TRAIN_FILE_PREFIX = 'Training_demo'\n",
        "TEST_FILE_PREFIX = 'Testing_demo'\n",
        "file_extension = '.tfrecord.gz'\n",
        "TRAIN_FILE_PATH = 'gs://' + DATA_BUCKET + '/' + TRAIN_FILE_PREFIX + file_extension\n",
        "TEST_FILE_PATH = 'gs://' + DATA_BUCKET + '/' + TEST_FILE_PREFIX + file_extension\n",
        "\n",
        "# The labels, consecutive integer indices starting from zero, are stored in\n",
        "# this property, set on each point.\n",
        "LABEL = 'landcover'\n",
        "# Number of label values, i.e. number of classes in the classification.\n",
        "N_CLASSES = 3\n",
        "\n",
        "# Use Landsat 8 surface reflectance data for predictors.\n",
        "L8SR = ee.ImageCollection('LANDSAT/LC08/C02/T1_L2')\n",
        "# Use these bands for prediction.\n",
        "BANDS = ['SR_B2', 'SR_B3', 'SR_B4', 'SR_B5', 'SR_B6', 'SR_B7']\n",
        "\n",
        "# These names are used to specify properties in the export of\n",
        "# training/testing data and to define the mapping between names and data\n",
        "# when reading into TensorFlow datasets.\n",
        "FEATURE_NAMES = BANDS + [LABEL]\n",
        "\n",
        "# List of fixed-length features, all of which are float32.\n",
        "columns = [\n",
        "  tf.io.FixedLenFeature(shape=[1], dtype=tf.float32) for k in FEATURE_NAMES\n",
        "]\n",
        "\n",
        "# Dictionary with feature names as keys, fixed-length features as values.\n",
        "FEATURES_DICT = dict(zip(FEATURE_NAMES, columns))\n",
        "\n",
        "ATTRIBUTION = 'Map Data \u0026copy; \u003ca href=\"https://earthengine.google.com/\"\u003eGoogle Earth Engine\u003c/a\u003e'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8sosRFEDdOMA"
      },
      "source": [
        "# Read data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43-c0JNFI_m6"
      },
      "source": [
        "### Check existence of the data files\n",
        "\n",
        "Check that you have permission to read the files in the output Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YDZfNl6yc0Kj"
      },
      "source": [
        "print('Found training file.' if tf.io.gfile.exists(TRAIN_FILE_PATH)\n",
        "    else 'No training file found.')\n",
        "print('Found testing file.' if tf.io.gfile.exists(TEST_FILE_PATH)\n",
        "    else 'No testing file found.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LS4jGTrEfz-1"
      },
      "source": [
        "## Read into a `tf.data.Dataset`\n",
        "\n",
        "Here we are going to read a file in Cloud Storage into a `tf.data.Dataset`.  ([these TensorFlow docs](https://www.tensorflow.org/guide/data) explain more about reading data into a `tf.data.Dataset`).  Check that you can read examples from the file.  The purpose here is to ensure that we can read from the file without an error.  The actual content is not necessarily human readable.  Note that we will use all data for training.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T3PKyDQW8Vpx",
        "cellView": "code"
      },
      "source": [
        "# Create a dataset from the TFRecord file in Cloud Storage.\n",
        "train_dataset = tf.data.TFRecordDataset(\n",
        "    [TRAIN_FILE_PATH, TEST_FILE_PATH], compression_type='GZIP')\n",
        "\n",
        "# Print the first record to check.\n",
        "print(iter(train_dataset).next())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QNfaUPbcjuCO"
      },
      "source": [
        "## Parse the dataset\n",
        "\n",
        "Now we need to make a parsing function for the data in the TFRecord files.  The data comes in flattened 2D arrays per record and we want to use the first part of the array for input to the model and the last element of the array as the class label.  The parsing function reads data from a serialized `Example` proto (i.e. [`example.proto`](https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/core/example/example.proto)) into a dictionary in which the keys are the feature names and the values are the tensors storing the value of the features for that example."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x2Q0g3fBj2kD",
        "cellView": "code"
      },
      "source": [
        "def parse_tfrecord(example_proto):\n",
        "  \"\"\"Deserialize an example proto.  Reshape to (1, 1, 1).\"\"\"\n",
        "  parsed_features = tf.io.parse_example(example_proto, FEATURES_DICT)\n",
        "  labels = parsed_features.pop(LABEL)\n",
        "  # (1) -\u003e (1, 1, 1)\n",
        "  return ({k: [[v]] for k, v in parsed_features.items()},\n",
        "          tf.expand_dims(tf.expand_dims(\n",
        "              tf.cast(labels, tf.int32), axis=0), axis=0))\n",
        "\n",
        "# Map the function over the dataset.\n",
        "parsed_dataset = train_dataset.map(parse_tfrecord, num_parallel_calls=4)\n",
        "\n",
        "# Print the first parsed record to check.\n",
        "check = iter(parsed_dataset.take(1)).next()\n",
        "pprint(check)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nEx1RAXOZQkS"
      },
      "source": [
        "# Model setup\n",
        "\n",
        "Make a densely-connected convolutional model, where the convolution occurs in a 1x1 kernel.  This allows Earth Engine to apply the model spatially, as demonstrated below.  Note that the model used here is purely for demonstration purposes."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t9pWa54oG-xl"
      },
      "source": [
        "## Create the Keras model\n",
        "This is a very basic [Keras functional model](https://www.tensorflow.org/guide/keras/functional)."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Note that the input is a [1, 1, P] vector of P bands, i.e. a 1x1 kernel.\n",
        "inputs = keras.Input(shape=(1, 1, len(BANDS)))\n",
        "x = tf.keras.layers.Conv2D(64, (1, 1), activation=tf.nn.relu)(inputs)\n",
        "x = tf.keras.layers.Dropout(0.1)(x)\n",
        "x = tf.keras.layers.Conv2D(N_CLASSES, (1, 1), activation=tf.nn.softmax)(x)\n",
        "model = tf.keras.Model(inputs=inputs, outputs=x)\n",
        "\n",
        "# A Layer to stack and reshape the input tensors.\n",
        "class MyPreprocessing(keras.layers.Layer):\n",
        "  def __init__(self, **kwargs):\n",
        "    super(MyPreprocessing, self).__init__(**kwargs)\n",
        "\n",
        "  def call(self, features_dict):\n",
        "    # (None, 1, 1, 1) -\u003e (None, 1, 1, P)\n",
        "    return tf.concat([features_dict[b] for b in BANDS], axis=3)\n",
        "\n",
        "  def get_config(self):\n",
        "    config = super().get_config()\n",
        "    return config\n",
        "\n",
        "\n",
        "# A Model that wraps the base model with the preprocessing layer.\n",
        "class MyModel(keras.Model):\n",
        "  def __init__(self, preprocessing, backbone, **kwargs):\n",
        "    super().__init__(**kwargs)\n",
        "    self.preprocessing = preprocessing\n",
        "    self.backbone = backbone\n",
        "\n",
        "  def call(self, features_dict):\n",
        "    x = self.preprocessing(features_dict)\n",
        "    return self.backbone(x)\n",
        "\n",
        "  def get_config(self):\n",
        "    config = super().get_config()\n",
        "    return config\n",
        "\n",
        "wrapped_model = MyModel(MyPreprocessing(), model)"
      ],
      "metadata": {
        "id": "sU2CO68NaJLR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OCZq3VNpG--G",
        "cellView": "code"
      },
      "source": [
        "# Compile the model with the specified loss and optimizer functions.\n",
        "wrapped_model.compile(\n",
        "    optimizer='adam',\n",
        "    loss='sparse_categorical_crossentropy',\n",
        "    metrics=['accuracy'])\n",
        "\n",
        "# Fit the model to the training data.  Lucky number 13.\n",
        "wrapped_model.fit(parsed_dataset.batch(4), epochs=13)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(wrapped_model.summary())"
      ],
      "metadata": {
        "id": "uzadZKRr_XHN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "keijPyVQTIAq"
      },
      "source": [
        "## De/serialization\n",
        "\n",
        "De/serialization prepares the model for hosting on Google Cloud.  Specifically, you need to provide image data to the Vertex AI API by sending the image data as base64-encoded text ([reference](https://cloud.google.com/vertex-ai/docs/general/base64)).  Wrap the trained model in de/serialization layers to handle the conversion."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class DeSerializeInput(tf.keras.layers.Layer):\n",
        "  def __init__(self, **kwargs):\n",
        "    super().__init__(**kwargs)\n",
        "\n",
        "  def call(self, inputs_dict):\n",
        "    return {\n",
        "      k: tf.map_fn(lambda x: tf.io.parse_tensor(x, tf.float32),\n",
        "                   tf.io.decode_base64(v),\n",
        "                   fn_output_signature=tf.float32)\n",
        "        for (k, v) in inputs_dict.items()\n",
        "    }\n",
        "\n",
        "  def get_config(self):\n",
        "    config = super().get_config()\n",
        "    return config\n",
        "\n",
        "\n",
        "class ReSerializeOutput(tf.keras.layers.Layer):\n",
        "  def __init__(self, **kwargs):\n",
        "    super().__init__(**kwargs)\n",
        "\n",
        "  def call(self, output_tensor):\n",
        "    return tf.map_fn(lambda x: tf.io.encode_base64(tf.io.serialize_tensor(x)),\n",
        "                    output_tensor,\n",
        "                    fn_output_signature=tf.string)\n",
        "\n",
        "  def get_config(self):\n",
        "    config = super().get_config()\n",
        "    return config\n",
        "\n",
        "input_deserializer = DeSerializeInput()\n",
        "output_deserializer = ReSerializeOutput()\n",
        "\n",
        "serialized_inputs = {\n",
        "    b: tf.keras.Input(shape=[], dtype='string', name=b) for b in BANDS\n",
        "}\n",
        "\n",
        "updated_model_input = input_deserializer(serialized_inputs)\n",
        "updated_model = wrapped_model(updated_model_input)\n",
        "updated_model = output_deserializer(updated_model)\n",
        "updated_model= tf.keras.Model(serialized_inputs, updated_model)"
      ],
      "metadata": {
        "id": "s6S2OoI1QCYb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tf.keras.utils.plot_model(updated_model)"
      ],
      "metadata": {
        "id": "palnJ5LgQpAH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "shbr6cSXShRg"
      },
      "source": [
        "## Save the trained model\n",
        "\n",
        "Export the trained model to TensorFlow `SavedModel` format in your cloud storage bucket.  The [Cloud Platform storage browser](https://console.cloud.google.com/storage/browser) is useful for checking on these saved models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sgg7MTXfS1PK"
      },
      "source": [
        "!gsutil rm -r {MODEL_DIR + '_test'}\n",
        "updated_model.save(MODEL_DIR + '_test')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0_uTqQAaVTIK"
      },
      "source": [
        "# Deploy the model on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### [Upload the model](https://cloud.google.com/sdk/gcloud/reference/ai/models/upload)\n",
        "Add an entry to the model registry that points to the location of the saved model and a container image needed to run the model."
      ],
      "metadata": {
        "id": "SNlHMjrin51r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai models delete {MODEL_NAME + '_test'} --project={PROJECT} --region={REGION}"
      ],
      "metadata": {
        "id": "rsUq4i8Redmk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8RZRRzcfVu5T"
      },
      "source": [
        "!gcloud ai models upload \\\n",
        "  --artifact-uri={MODEL_DIR + '_test'} \\\n",
        "  --project={PROJECT} \\\n",
        "  --region={REGION} \\\n",
        "  --container-image-uri={CONTAINER_IMAGE} \\\n",
        "  --description={MODEL_NAME + '_test'} \\\n",
        "  --display-name={MODEL_NAME + '_test'} \\\n",
        "  --model-id={MODEL_NAME + '_test'}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### [Create an endpoint](https://cloud.google.com/sdk/gcloud/reference/ai/endpoints/create)\n",
        "\n",
        "Create an endpoint from which to serve the model."
      ],
      "metadata": {
        "id": "zwaYjMgYoHD6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai endpoints create \\\n",
        "  --display-name={ENDPOINT_NAME + '_test'} \\\n",
        "  --region={REGION} \\\n",
        "  --project={PROJECT}"
      ],
      "metadata": {
        "id": "7tJZt_XcWmr7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### [Deploy the model to the endpoint](https://cloud.google.com/sdk/gcloud/reference/ai/endpoints/deploy-model)\n",
        "\n",
        "First, look up the endpoint ID, then deploy the model."
      ],
      "metadata": {
        "id": "QS_Rksv-wNQg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ENDPOINT_ID = !gcloud ai endpoints list \\\n",
        "  --project={PROJECT} \\\n",
        "  --region={REGION} \\\n",
        "  --filter=displayName:{ENDPOINT_NAME + '_test'} \\\n",
        "  --format=\"value(ENDPOINT_ID.scope())\"\n",
        "ENDPOINT_ID = ENDPOINT_ID[-1]"
      ],
      "metadata": {
        "id": "m3tgIGXld4xw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai endpoints deploy-model {ENDPOINT_ID} \\\n",
        "  --project={PROJECT} \\\n",
        "  --region={REGION} \\\n",
        "  --model={MODEL_NAME + '_test'} \\\n",
        "  --display-name={MODEL_NAME + '_test'}"
      ],
      "metadata": {
        "id": "TIO93cnuwXhb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5aTGza-rWIjp"
      },
      "source": [
        "# Connect to the hosted model from Earth Engine\n",
        "\n",
        "1. Generate the input imagery.  This should be done in exactly the same way as the training data were generated.  See [this example notebook](http://colab.research.google.com/github/google/earthengine-api/blob/master/python/examples/ipynb/TF_demo1_keras.ipynb) for details.\n",
        "2. Connect to the hosted model.\n",
        "3. Use the model to make predictions.\n",
        "4. Display the results.\n",
        "\n",
        "Note that it may take the model a couple minutes to spin up before it is ready to make predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P2OsyrJ7HAhE"
      },
      "source": [
        "def maskL8sr(image):\n",
        "  \"\"\"Cloud masking function for Landsat 8, collection 2.\"\"\"\n",
        "  qa_mask = image.select('QA_PIXEL').bitwiseAnd(31).eq(0)\n",
        "  saturation_mask = image.select('QA_RADSAT').eq(0)\n",
        "\n",
        "  optical_bands = image.select('SR_B.').multiply(0.0000275).add(-0.2)\n",
        "  thermal_bands = image.select('ST_B.*').multiply(0.00341802).add(149.0)\n",
        "\n",
        "  return (image.addBands(optical_bands, None, True)\n",
        "                .addBands(thermal_bands, None, True)\n",
        "                .updateMask(qa_mask)\n",
        "                .updateMask(saturation_mask).select('SR_B*.'))\n",
        "\n",
        "\n",
        "# The image input data is a 2018 cloud-masked median composite.\n",
        "image = L8SR.filterDate('2018-01-01', '2018-12-31').map(maskL8sr).median()\n",
        "\n",
        "# Get a URL to serve image tiles.\n",
        "mapid = image.getMapId({'bands': ['SR_B4', 'SR_B3', 'SR_B2'], 'min': 0, 'max': 0.3})\n",
        "\n",
        "# Use folium to visualize the imagery.\n",
        "map = folium.Map(location=[38., -122.5])\n",
        "\n",
        "# Inputs.\n",
        "folium.TileLayer(\n",
        "    tiles=mapid['tile_fetcher'].url_format,\n",
        "    attr=ATTRIBUTION,\n",
        "    overlay=True,\n",
        "    name='median composite',\n",
        "  ).add_to(map)\n",
        "\n",
        "endpoint_path = (\n",
        "    'projects/' + PROJECT + '/locations/' + REGION + '/endpoints/' + str(ENDPOINT_ID))\n",
        "\n",
        "# Connect to the hosted model.\n",
        "vertex_model = ee.Model.fromVertexAi(\n",
        "  endpoint=endpoint_path,\n",
        "  inputTileSize=[8, 8],\n",
        "  proj=ee.Projection('EPSG:4326').atScale(30),\n",
        "  fixInputProj=True,\n",
        "  outputBands={'output': {\n",
        "      'type': ee.PixelType.float(),\n",
        "      'dimensions': 1\n",
        "    }\n",
        "  })\n",
        "\n",
        "predictions = vertex_model.predictImage(image.select(BANDS).float())\n",
        "probabilities = predictions.arrayFlatten([['bare', 'veg', 'water']])\n",
        "probability_vis = {'bands': ['bare', 'veg', 'water'], 'min': 0.2, 'max': 0.5}\n",
        "probability_mapid = probabilities.getMapId(probability_vis)\n",
        "folium.TileLayer(\n",
        "    tiles=probability_mapid['tile_fetcher'].url_format,\n",
        "    attr=ATTRIBUTION,\n",
        "    overlay=True,\n",
        "    name='predictions',\n",
        "  ).add_to(map)\n",
        "\n",
        "map.add_child(folium.LayerControl())\n",
        "display(map)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}