guides/linked/Earth_Engine_TensorFlow_tree_counting_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "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_tree_counting_model.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_tree_counting_model.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",
      "source": [
        "# Counting trees with Earth Engine and a custom TensorFlow model\n",
        "\n",
        "Ever wonder how many trees there are?  Of course you have.  We were wondering the same thing.  People have been counting trees in imagery for a good long time ([Wang et al. 2004](https://doi.org/10.14358/PERS.70.3.351) have a nice little summary).  This kind of counting is useful for forest inventory, to understand species composition, density, diameter and height distributions.  Given that convolutional neural networks are fashionable nowadays, we wondered if tree-crown segmentation was a thing.  It is!  [Li et al. 2023](https://doi.org/10.1093/pnasnexus/pgad076) have a great paper, just out.  It's also commendable that they publish their code and other artifacts on GitHub ([their repo](https://github.com/sizhuoli/TreeCountSegHeight/tree/main)).  Specifically, they've got a bunch of trained TensorFlow models in there.  Jackpot!"
      ],
      "metadata": {
        "id": "jm0ZLV8YRb-r"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Host a trained model on Vertex AI to get predictions in Earth Engine\n",
        "\n",
        "If there's suitable input data in Earth Engine (check the [extensive data catalog](https://developers.google.com/earth-engine/datasets) or [upload your own](https://developers.google.com/earth-engine/guides/image_upload)), you can get interactive predictions from geospatially-aware models hosted on Vertex AI.  By geospatially aware, we mean models that accept three-dimensional (four if you include the batch dimension) inputs: `[height, width, channels]`.  For deployment in a map, you also need to know the scale of the pixel, in ground units (meters).  [Li et al. 2023](https://doi.org/10.1093/pnasnexus/pgad076) explain most of this, and the rest is discoverable from the saved model itself.\n",
        "\n",
        "To get started, import the necessary libraries and authenticate.\n",
        "\n",
        "#### **Warning!** This demo consumes billable resources of Google Cloud, including Earth Engine, Vertex AI and Cloud Storage."
      ],
      "metadata": {
        "id": "g2OWp508XO7c"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JtqbQ2K20UMu"
      },
      "outputs": [],
      "source": [
        "import ee\n",
        "import folium\n",
        "import google\n",
        "import tensorflow as tf\n",
        "\n",
        "from google.colab import auth\n",
        "from keras import backend as K\n",
        "from oauth2client.client import GoogleCredentials\n",
        "from pydrive.auth import GoogleAuth\n",
        "from pydrive.drive import GoogleDrive"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "auth.authenticate_user()\n",
        "credentials, project = google.auth.default()"
      ],
      "metadata": {
        "id": "HXfFSx4P0rU9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# REPLACE WITH YOUR CLOUD PROJECT!\n",
        "PROJECT = 'your-project'\n",
        "# REPLACE WITH YOUR (WRITABLE) CLOUD BUCKET!\n",
        "BUCKET = 'your-bucket'\n",
        "\n",
        "# Hosted model and endpoint names.\n",
        "MODEL_NAME = 'trees_20210620-0202_adam_e4_redgreenblue_256'\n",
        "ENDPOINT_NAME = 'trees_endpoint'\n",
        "REGION = 'us-central1'\n",
        "\n",
        "# A container image that can run the hosted model.\n",
        "# See https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers#tensorflow\n",
        "CONTAINER_IMAGE = 'us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-11:latest'\n",
        "\n",
        "ATTRIBUTION = 'Map Data \u0026copy; \u003ca href=\"https://earthengine.google.com/\"\u003eGoogle Earth Engine\u003c/a\u003e'"
      ],
      "metadata": {
        "id": "RQkql50sfId1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nqQ4lkyPNKuY"
      },
      "outputs": [],
      "source": [
        "ee.Initialize(credentials, project=PROJECT)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download a trained model from GitHub\n",
        "\n",
        "Grab a set of trained models from the link in the [GitHub repo](https://github.com/sizhuoli/TreeCountSegHeight/blob/main/models/readme_models.md).  Note that there are several models, with different input set options.  We're assuming that all we have is RGB imagery (i.e. no NIR or DSM data).  Transfer that model to your Cloud Storage bucket.  From there, load the model directly.  Note that we're passing the Keras backend as a custom object.  Since we don't need to (re)train the model, we don't need to compile it."
      ],
      "metadata": {
        "id": "-28i_ep2ZgwV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "gauth = GoogleAuth()\n",
        "gauth.credentials = GoogleCredentials.get_application_default()\n",
        "drive = GoogleDrive(gauth)\n",
        "\n",
        "file_id = '1ZNibrh6pa4-cjXLawua6L96fOKS3uwbn'\n",
        "downloaded = drive.CreateFile({'id': file_id})\n",
        "downloaded.GetContentFile('saved_models.zip')"
      ],
      "metadata": {
        "id": "2cHbCo5093-I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!unzip saved_models.zip"
      ],
      "metadata": {
        "id": "Vziz7bdFAFCp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls -l saved_models"
      ],
      "metadata": {
        "id": "oG8tzYpZB0HQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Grab the RGB model.\n",
        "RGB_MODEL = 'trees_20210620-0202_Adam_e4_redgreenblue_256_84_frames_weightmapTversky_MSE100_5weight_attUNet.h5'"
      ],
      "metadata": {
        "id": "wI35wHtxCFvv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_file = 'saved_models/' + RGB_MODEL\n",
        "model = tf.keras.models.load_model(model_file, custom_objects={'K': K}, compile=False)"
      ],
      "metadata": {
        "id": "5Csqk3pw1m0V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run a few checks to investigate the model."
      ],
      "metadata": {
        "id": "mtunQ_xlbN0P"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "if model.trainable:\n",
        "  print('The model is trained.')\n",
        "else:\n",
        "  print('The model is not trained.')"
      ],
      "metadata": {
        "id": "xuRdh6Yg9MeB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.summary()"
      ],
      "metadata": {
        "id": "9raWBwMT5wd4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "zeros = tf.zeros((1, 256, 256, 3))\n",
        "model(zeros)"
      ],
      "metadata": {
        "id": "gm6MPSPs53vt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Prepare the model for hosting\n",
        "\n",
        "Note that there is a list of two outputs.  Layers called `output_seg` and `output_dens`.  Wrap the model in de/serialization layers, for hosting on Vertex AI.  Save the wrapped model."
      ],
      "metadata": {
        "id": "fOpuYBwFbUvA"
      }
    },
    {
      "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_dict={}\n",
        "\n",
        "    for (k, v) in inputs_dict.items():\n",
        "      decoded = tf.io.decode_base64(v)\n",
        "      return_dict[k] = tf.map_fn(lambda x: tf.io.parse_tensor(x, tf.float32),\n",
        "                                 decoded, fn_output_signature=tf.float32)\n",
        "\n",
        "    return return_dict\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, outputs_list):\n",
        "      return [tf.map_fn(\n",
        "          lambda x: tf.io.encode_base64(\n",
        "              tf.io.serialize_tensor(x)), tensor, fn_output_signature=tf.string)\n",
        "      for tensor in outputs_list]\n",
        "\n",
        "  def get_config(self):\n",
        "    config = super().get_config()\n",
        "    return config\n",
        "\n",
        "\n",
        "input_deserializer = DeSerializeInput()\n",
        "output_deserializer = ReSerializeOutput()\n",
        "\n",
        "serialized_inputs = {\n",
        "    model.inputs[0].name: tf.keras.Input(shape=[], dtype='string', name='array_image')\n",
        "}\n",
        "\n",
        "updated_model_input = input_deserializer(serialized_inputs)\n",
        "updated_model = model(updated_model_input)\n",
        "updated_model = output_deserializer(updated_model)\n",
        "updated_model = tf.keras.Model(serialized_inputs, updated_model)"
      ],
      "metadata": {
        "id": "CJFeloY56HPv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "MODEL_DIR = f'gs://{BUCKET}/{RGB_MODEL[:-3]}'\n",
        "updated_model.save(MODEL_DIR)"
      ],
      "metadata": {
        "id": "IheGsouU8Nal"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Host the model\n",
        "\n",
        "First delete any models of the same name.  If you get an error that model doesn't exist, you can ignore that.  If you get an error because the model exists and is deployed to an endpoint, you will need to either rename the model or undeploy the previously deployed model from the [Cloud Console](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints)."
      ],
      "metadata": {
        "id": "CZR_xuEdcNpv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai models delete {MODEL_NAME} --project={PROJECT} --region={REGION}"
      ],
      "metadata": {
        "id": "nMLcryALLUdN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai models upload \\\n",
        "  --artifact-uri={MODEL_DIR} \\\n",
        "  --project={PROJECT} \\\n",
        "  --region={REGION} \\\n",
        "  --container-image-uri={CONTAINER_IMAGE} \\\n",
        "  --description={MODEL_NAME} \\\n",
        "  --display-name={MODEL_NAME} \\\n",
        "  --model-id={MODEL_NAME}"
      ],
      "metadata": {
        "id": "ua-X7QXc97c4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai endpoints create \\\n",
        "  --display-name={ENDPOINT_NAME} \\\n",
        "  --region={REGION} \\\n",
        "  --project={PROJECT}"
      ],
      "metadata": {
        "id": "lGWvtmFb-6He"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ENDPOINT_ID = !gcloud ai endpoints list \\\n",
        "  --project={PROJECT} \\\n",
        "  --region={REGION} \\\n",
        "  --filter=displayName:{ENDPOINT_NAME} \\\n",
        "  --format=\"value(ENDPOINT_ID.scope())\"\n",
        "ENDPOINT_ID = ENDPOINT_ID[-1]"
      ],
      "metadata": {
        "id": "kSsrmjsM-1aq"
      },
      "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} \\\n",
        "  --machine-type=n1-standard-8 \\\n",
        "  --accelerator=type=nvidia-tesla-t4,count=1 \\\n",
        "  --display-name={MODEL_NAME}"
      ],
      "metadata": {
        "id": "SmgjnrQj-1ar"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5aTGza-rWIjp"
      },
      "source": [
        "## Connect to the hosted model from Earth Engine\n",
        "\n",
        "It may take a few minutes for the model to display predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P2OsyrJ7HAhE"
      },
      "source": [
        "estonia_images = ee.ImageCollection(\"Estonia/Maamet/orthos/rgb\")\n",
        "image = estonia_images.mosaic().float()\n",
        "\n",
        "# Get a URL to serve image tiles.\n",
        "mapid = image.getMapId({'bands': ['R', 'G', 'B'], 'min': 0, 'max': 128})\n",
        "\n",
        "# Use folium to visualize the imagery.\n",
        "map = folium.Map(location=[59.246333, 25.463968], zoom_start=21)\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/' + ENDPOINT_ID)\n",
        "\n",
        "# Connect to the hosted model.\n",
        "vertex_model = ee.Model.fromVertexAi(\n",
        "  endpoint=endpoint_path,\n",
        "  inputTileSize=[128, 128],\n",
        "  inputOverlapSize=[64, 64],\n",
        "  proj=ee.Projection('EPSG:4326').atScale(0.2),\n",
        "  fixInputProj=True,\n",
        "  outputBands={\n",
        "    're_serialize_output': {\n",
        "      'type': ee.PixelType.float(),\n",
        "      'dimensions': 1\n",
        "    },\n",
        "    're_serialize_output_1': {\n",
        "      'type': ee.PixelType.float(),\n",
        "      'dimensions': 1\n",
        "    }\n",
        "  }\n",
        ")\n",
        "\n",
        "# Normalization.\n",
        "stats = image.reduceNeighborhood(\n",
        "  reducer=ee.Reducer.mean().combine(ee.Reducer.stdDev(), None, True),\n",
        "  kernel=ee.Kernel.square(10, 'meters'),\n",
        "  optimization='window'\n",
        ")\n",
        "means = stats.select(['R_mean', 'G_mean', 'B_mean'])\n",
        "sds = stats.select(['R_stdDev', 'G_stdDev', 'B_stdDev'])\n",
        "input_image = (image.select(['R', 'G', 'B'])\n",
        "    .subtract(means).divide(sds).float().toArray().rename('array_image'))\n",
        "\n",
        "# Predictions.\n",
        "predictions = vertex_model.predictImage(input_image)\n",
        "seg = predictions.select('re_serialize_output').arrayGet([0])\n",
        "prob = predictions.select('re_serialize_output_1').arrayGet([0])\n",
        "\n",
        "seg_mapid = seg.getMapId({'min': 0, 'max': 0.5})\n",
        "folium.TileLayer(\n",
        "    tiles=seg_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": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### **Warning!** This demo consumes billable resources of Google Cloud, including Earth Engine, Vertex AI and Cloud Storage.  Be sure to shut down any prediction nodes to avoid ongoing charges."
      ],
      "metadata": {
        "id": "67KfBl7z2OSX"
      }
    }
  ]
}