guides/linked/Earth_Engine_TensorFlow_DNN_from_scratch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "private_outputs": true,
      "gpuType": "T4",
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "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_DNN_from_scratch.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_DNN_from_scratch.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": [
        "# Introduction\n",
        "\n",
        "This is an Earth Engine and TensorFlow demonstration notebook.  Specifically, this notebook shows:\n",
        "\n",
        "1.   Exporting training/testing data from Earth Engine in TFRecord format.\n",
        "2.   Preparing the data for use in a TensorFlow model.\n",
        "2.   Training and validating a simple model in TensorFlow.\n",
        "3.   Making predictions on image data exported from Earth Engine in TFRecord format.\n",
        "4.   Ingesting classified image data to Earth Engine in TFRecord format.\n",
        "\n",
        "This is intended to demonstrate a complete i/o pipeline.  For a workflow that uses a [Vertex AI Platform](https://cloud.google.com/vertex-ai) hosted model making predictions interactively, see [this example notebook](http://colab.research.google.com/github/google/earthengine-api/blob/master/python/examples/ipynb/Earth_Engine_TensorFlow_Vertex_AI.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KiTyR3FNlv-O"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import auth\n",
        "from pprint import pprint\n",
        "import ee\n",
        "import folium\n",
        "import google\n",
        "import json\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras"
      ],
      "metadata": {
        "id": "Ai7UHDxc9itO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DrXLkJC2QJdP"
      },
      "source": [
        "# Define variables\n",
        "\n",
        "This set of global variables will be used throughout.  For this demo, you must have a Cloud Storage bucket into which you can write files."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GHTOc5YLQZ5B"
      },
      "source": [
        "# REPLACE WITH YOUR CLOUD PROJECT\n",
        "PROJECT = 'your-project'\n",
        "\n",
        "# REPLACE WITH YOUR EARTH ENGINE ASSETS PROJECT\n",
        "ASSETS_PROJECT = 'your-ee-project'\n",
        "\n",
        "# Cloud Storage bucket into which training, testing and prediction\n",
        "# datasets will be written.  You must be able to write into this bucket.\n",
        "# REPLACE WITH YOUR CLOUD STORAGE BUCKET\n",
        "OUTPUT_BUCKET = 'your-bucket'\n",
        "\n",
        "# Use Landsat 8 surface reflectance bands for predictors.\n",
        "BANDS = ['SR_B2', 'SR_B3', 'SR_B4', 'SR_B5', 'SR_B6', 'SR_B7']\n",
        "INPUTS = BANDS + ['NDVI']\n",
        "\n",
        "# The labels, consecutive integer indices starting from zero, are stored in\n",
        "# this property, set on each point.\n",
        "LABEL = 'landcover'\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",
        "# Number of label values, i.e. number of classes in the classification.\n",
        "N_CLASSES = 3\n",
        "\n",
        "# Pixel scale in meters.\n",
        "SCALE = 30\n",
        "\n",
        "# File names for the training and testing datasets.  These TFRecord files\n",
        "# will be exported from Earth Engine into 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://' + OUTPUT_BUCKET + '/' + TRAIN_FILE_PREFIX + file_extension\n",
        "TEST_FILE_PATH = 'gs://' + OUTPUT_BUCKET + '/' + TEST_FILE_PREFIX + file_extension\n",
        "\n",
        "# File name for the prediction (image) dataset.  The trained model will read\n",
        "# this dataset and make predictions in each pixel.\n",
        "PREDICTION_INPUT_PREFIX = 'Image_pixel_demo_'\n",
        "\n",
        "# The output path for the classified image (i.e. predictions) TFRecord file.\n",
        "PREDICTION_OUTPUT_FILE = f'gs://{OUTPUT_BUCKET}/Classified_pixel_demo.TFRecord'\n",
        "\n",
        "# Set of min-max longitudes and latitudes.  Export imagery in this region.\n",
        "PREDICTION_REGION_COORDS = [-122.7, 37.3, -121.8, 38.00]\n",
        "\n",
        "# The name of the Earth Engine asset to be created by importing\n",
        "# the classified image from the TFRecord file in Cloud Storage.\n",
        "PREDICTION_ASSET_ID = f'projects/{ASSETS_PROJECT}/assets/DNN_from_scratch_demo'\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",
      "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": [
        "credentials, project_id = 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": "ZcjQnHH8zT4q"
      },
      "source": [
        "# Get Training and Testing data from Earth Engine\n",
        "\n",
        "To get data for a classification model of three classes (bare, vegetation, water), we need labels and the value of predictor variables for each labeled example.  We've already generated some labels in Earth Engine.  Specifically, these are visually interpreted points labeled \"bare,\" \"vegetation,\" or \"water\" for a very simple classification demo ([example script](https://code.earthengine.google.com/?scriptPath=Examples%3ADemos%2FClassification)).  For predictor variables, we'll use [Landsat 8 surface reflectance imagery](https://developers.google.com/earth-engine/datasets/catalog/LANDSAT_LC08_C02_T1_L2), bands 2-7."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0EJfjgelSOpN"
      },
      "source": [
        "## Prepare Landsat 8 imagery\n",
        "\n",
        "First, make a cloud-masked median composite of Landsat 8 surface reflectance imagery from 2018.  Check the composite by visualizing with folium."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DJYucYe3SPPr"
      },
      "source": [
        "def maskL8sr(image):\n",
        "  \"\"\"Cloud masking function for Landsat 8 surface reflectance.\"\"\"\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",
        "composite = (ee.ImageCollection('LANDSAT/LC08/C02/T1_L2')\n",
        "          .filterDate('2018-01-01', '2018-12-31').map(maskL8sr).median())\n",
        "\n",
        "# Use folium to visualize the imagery.\n",
        "mapid = composite.getMapId({'bands': ['SR_B4', 'SR_B3', 'SR_B2'], 'min': 0, 'max': 0.3})\n",
        "map = folium.Map(location=[38., -122.5])\n",
        "\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",
        "map.add_child(folium.LayerControl())\n",
        "map"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UEeyPf3zSPct"
      },
      "source": [
        "## Add pixel values of the composite to labeled points\n",
        "\n",
        "Some training labels have already been collected for you.  Load the labeled points from an existing Earth Engine asset.  Each point in this table has a property called `landcover` that stores the label, encoded as an integer.  Here we overlay the points on  imagery to get predictor variables along with labels."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iOedOKyRExHE"
      },
      "source": [
        "# The training/testing is a dataset of points with known land cover labels.\n",
        "# Sample the image at the points and add a random column.\n",
        "sample = composite.sampleRegions(\n",
        "  collection=ee.FeatureCollection('projects/google/demo_landcover_labels'),\n",
        "  properties=[LABEL],\n",
        "  scale=SCALE).randomColumn()\n",
        "\n",
        "# Partition the sample approximately 70-30.\n",
        "training = sample.filter(ee.Filter.lt('random', 0.7))\n",
        "testing = sample.filter(ee.Filter.gte('random', 0.7))\n",
        "\n",
        "# Print the first couple points to verify.\n",
        "pprint({'training': training.first().getInfo()})\n",
        "pprint({'testing': testing.first().getInfo()})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uNc7a2nRR4MI"
      },
      "source": [
        "## Export the training and testing data\n",
        "\n",
        "Now that there's training and testing data in Earth Engine and you've inspected a couple examples to ensure that the information you need is present, it's time to materialize the datasets in a place where the TensorFlow model has access to them.  You can do that by exporting the training and testing datasets to tables in TFRecord format ([learn more about TFRecord format](https://www.tensorflow.org/tutorials/load_data/tfrecord)) in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pb-aPvQc0Xvp"
      },
      "source": [
        "# Make sure you can see the output bucket.  You must have write access.\n",
        "print('Found Cloud Storage bucket.' if tf.io.gfile.exists('gs://' + OUTPUT_BUCKET)\n",
        "    else 'Can not find output Cloud Storage bucket.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wtoqj0Db1TmJ"
      },
      "source": [
        "Once you've verified the existence of the intended output bucket, run the exports."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TfVNQzg8R6Wy"
      },
      "source": [
        "# Create the tasks.\n",
        "training_task = ee.batch.Export.table.toCloudStorage(\n",
        "  collection=training,\n",
        "  description='Training Export',\n",
        "  fileNamePrefix=TRAIN_FILE_PREFIX,\n",
        "  bucket=OUTPUT_BUCKET,\n",
        "  fileFormat='TFRecord',\n",
        "  selectors=FEATURE_NAMES)\n",
        "\n",
        "testing_task = ee.batch.Export.table.toCloudStorage(\n",
        "  collection=testing,\n",
        "  description='Testing Export',\n",
        "  fileNamePrefix=TEST_FILE_PREFIX,\n",
        "  bucket=OUTPUT_BUCKET,\n",
        "  fileFormat='TFRecord',\n",
        "  selectors=FEATURE_NAMES)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QF4WGIekaS2s"
      },
      "source": [
        "# Start the tasks.\n",
        "training_task.start()\n",
        "testing_task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q7nFLuySISeC"
      },
      "source": [
        "### Monitor task progress\n",
        "\n",
        "You can see all your Earth Engine tasks by listing them.  Make sure the training and testing tasks are completed before continuing."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oEWvS5ekcEq0"
      },
      "source": [
        "# Print all tasks.\n",
        "pprint(ee.batch.Task.list())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43-c0JNFI_m6"
      },
      "source": [
        "### Check existence of the exported files\n",
        "\n",
        "If you've seen the status of the export tasks change to `COMPLETED`, then check for the existence of 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": "NA8QA8oQVo8V"
      },
      "source": [
        "## Export the imagery\n",
        "\n",
        "Now that there's training and testing data, you can train a model, host the model on Vertex AI and get predictions directly in Earth Engine, as shown in [this demo](/earth-engine/guides/tf_examples#multi-class-prediction-with-a-dnn-hosted-on-vertex-ai).  Alternatively, you may want to export imagery, run prediction outside Earth Engine, then import the resultant predictions to Earth Engine as an image.  To continue with the latter approach, export whatever imagery on which you want to perform inference to the output Cloud Storage bucket as TFRecord files."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tVNhJYacVpEw"
      },
      "source": [
        "# Specify patch and file dimensions.\n",
        "image_export_options = {\n",
        "  'patchDimensions': [256, 256],\n",
        "  'maxFileSize': 104857600,\n",
        "  'compressed': True\n",
        "}\n",
        "\n",
        "# Setup the task.\n",
        "image_task = ee.batch.Export.image.toCloudStorage(\n",
        "  image=composite,\n",
        "  description='Export TFRecord imagery for inference demo',\n",
        "  fileNamePrefix=PREDICTION_INPUT_PREFIX,\n",
        "  bucket=OUTPUT_BUCKET,\n",
        "  scale=SCALE,\n",
        "  fileFormat='TFRecord',\n",
        "  region=ee.Geometry.Rectangle(PREDICTION_REGION_COORDS).getInfo()['coordinates'],\n",
        "  formatOptions=image_export_options,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6SweCkHDaNE3"
      },
      "source": [
        "# Start the task.\n",
        "image_task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JC8C53MRTG_E"
      },
      "source": [
        "### Monitor task progress"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BmPHb779KOXm"
      },
      "source": [
        "# Print all tasks.\n",
        "pprint(ee.batch.Task.list())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SrUhA1JKLONj"
      },
      "source": [
        "It's also possible to monitor an individual task.  Here we poll the task until it's done.  If you do this, please put a `sleep()` in the loop to avoid making too many requests.  Note that this will block until complete (you can always halt the execution of this cell)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rKZeZswloP11"
      },
      "source": [
        "import time\n",
        "\n",
        "while image_task.active():\n",
        "  print('Polling for task (id: {}).'.format(image_task.id))\n",
        "  time.sleep(30)\n",
        "print('Done with image export.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9vWdH_wlZCEk"
      },
      "source": [
        "# Data preparation and pre-processing\n",
        "\n",
        "Read data from the TFRecord file into a `tf.data.Dataset`.  Pre-process the dataset to get it into a suitable format for input to the model."
      ]
    },
    {
      "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 `Dataset`).\n",
        "\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(TRAIN_FILE_PATH, compression_type='GZIP')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ensure that you can read from the file without an error by getting a raw record.  Inspect the content of the record."
      ],
      "metadata": {
        "id": "U7UrS5hP4U6C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Get the first record to check.\n",
        "raw_record = iter(train_dataset).next()\n",
        "print('raw_record:')\n",
        "print(raw_record)\n",
        "\n",
        "# Get the raw record as a numpy array.\n",
        "record = raw_record.numpy()\n",
        "print('record:')\n",
        "print(record)\n",
        "\n",
        "# Decode a serialized tf.Example as a proto.\n",
        "example_proto = tf.train.Example.FromString(record)\n",
        "print('example_proto:')\n",
        "print(example_proto)"
      ],
      "metadata": {
        "id": "S7hTnP8B4KZ7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BrDYm-ibKR6t"
      },
      "source": [
        "## Define the structure of your data\n",
        "\n",
        "For parsing the exported TFRecord files, `features_dict` is a mapping between feature names (recall that `features_dict` contains the band and label names) and `float32` [`tf.io.FixedLenFeature`](https://www.tensorflow.org/api_docs/python/tf/io/FixedLenFeature) objects.  This mapping is necessary for telling TensorFlow how to read data in a TFRecord file into tensors.  Specifically, **all numeric data exported from Earth Engine is exported as `float32`**.\n",
        "\n",
        "(Note: *features* in the TensorFlow context (i.e. [`tf.train.Feature`](https://www.tensorflow.org/api_docs/python/tf/train/Feature)) are not to be confused with Earth Engine features (i.e. [`ee.Feature`](https://developers.google.com/earth-engine/apidocs/ee-feature)), where the former is a protocol message type for serialized data input to the model and the latter is a geometry-based geographic data structure.)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-6JVQV5HKHMZ",
        "cellView": "code"
      },
      "source": [
        "# List of fixed-length features, all of which are float32.\n",
        "columns = [\n",
        "  tf.io.FixedLenFeature(shape=(), dtype=tf.float32) for k in FEATURE_NAMES\n",
        "]\n",
        "\n",
        "# Dictionary with names as keys, features as values.\n",
        "features_dict = dict(zip(FEATURE_NAMES, columns))\n",
        "\n",
        "pprint(features_dict)"
      ],
      "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](https://www.tensorflow.org/api_docs/python/tf/train/Example) 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",
        "  \"\"\"The parsing function.\n",
        "\n",
        "  Read a serialized example into the structure defined by featuresDict.\n",
        "\n",
        "  Args:\n",
        "    example_proto: a serialized Example.\n",
        "\n",
        "  Returns:\n",
        "    A tuple of the predictors dictionary and the label, cast to an `int32`.\n",
        "  \"\"\"\n",
        "  parsed_features = tf.io.parse_example(example_proto, features_dict)\n",
        "  labels = parsed_features.pop(LABEL)\n",
        "  return parsed_features, tf.cast(labels, tf.int32)\n",
        "\n",
        "# Map the function over the dataset.\n",
        "parsed_dataset = train_dataset.map(parse_tfrecord, num_parallel_calls=5)\n",
        "\n",
        "# Print the first parsed record to check.\n",
        "pprint(iter(parsed_dataset).next())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nb8EyNT4Xnhb"
      },
      "source": [
        "Note that each record of the parsed dataset contains a tuple.  The first element of the tuple is a dictionary with bands for keys and the numeric value of the bands for values.  The second element of the tuple is a class label."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xLCsxWOuEBmE"
      },
      "source": [
        "## Create additional features\n",
        "\n",
        "Another thing we might want to do as part of the input process is to create new features, for example NDVI, a vegetation index computed from reflectance in two spectral bands.  Here are some helper functions for that."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lT6v2RM_EB1E",
        "cellView": "code"
      },
      "source": [
        "def normalized_difference(a, b):\n",
        "  \"\"\"Compute normalized difference of two inputs.\n",
        "\n",
        "  Compute (a - b) / (a + b).\n",
        "\n",
        "  Args:\n",
        "    a: an input tensor with shape=[1]\n",
        "    b: an input tensor with shape=[1]\n",
        "\n",
        "  Returns:\n",
        "    The normalized difference as a tensor.\n",
        "  \"\"\"\n",
        "  return tf.math.divide_no_nan((a - b), (a + b))\n",
        "\n",
        "def add_NDVI(features):\n",
        "  \"\"\"Add NDVI to the dataset.\n",
        "  Args:\n",
        "    features: a dictionary of input tensors keyed by feature name.\n",
        "    label: the target label\n",
        "\n",
        "  Returns:\n",
        "    A tuple of the input dictionary with an NDVI tensor added and the label.\n",
        "  \"\"\"\n",
        "  features['NDVI'] = normalized_difference(features['SR_B5'], features['SR_B4'])\n",
        "  return features"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nEx1RAXOZQkS"
      },
      "source": [
        "# Model setup\n",
        "\n",
        "The basic workflow for classification in TensorFlow is:\n",
        "\n",
        "1.  Create the model.\n",
        "2.  Train the model (i.e. `fit()`).\n",
        "3.  Use the trained model for inference (i.e. `predict()`).\n",
        "\n",
        "Here we'll create a neural network model using Keras.  Note that the model used here is purely for demonstration purposes and hasn't gone through any performance tuning."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t9pWa54oG-xl"
      },
      "source": [
        "## Create the Keras model\n",
        "\n",
        "Before we create the model, there's still a bit of pre-processing to get the data into the right input shape and a format that can be used with cross-entropy loss.  Specifically, Keras expects a list of inputs and a one-hot vector for the class. (See [the Keras loss function docs](https://keras.io/losses/), [the TensorFlow categorical identity docs](https://www.tensorflow.org/guide/feature_columns#categorical_identity_column) and [the `tf.one_hot` docs](https://www.tensorflow.org/api_docs/python/tf/one_hot) for details).\n",
        "\n",
        "Here we will use a simple neural network model with a 64 node hidden layer, a dropout layer and an output layer.  Use a custom pre-processing `Layer` to add an NDVI input feature and stack the inputs into a vector.  The following defines the pre-processing layer, wraps the model in the preprocessing layer, compiles it, and fits it to the training data."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# The base model takes a vector of inputs and returns a vector of outputs.\n",
        "inputs = keras.Input(shape=(len(INPUTS)), name='input_array')\n",
        "x = tf.keras.layers.Dense(64, activation=tf.nn.relu)(inputs)\n",
        "x = tf.keras.layers.Dropout(0.1)(x)\n",
        "x = tf.keras.layers.Dense(N_CLASSES, activation=tf.nn.softmax)(x)\n",
        "model = tf.keras.Model(inputs=inputs, outputs=x)\n",
        "\n",
        "# A Layer to add an NDVI feature and stack 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",
        "    features_dict = add_NDVI(features_dict)\n",
        "    return tf.stack([features_dict[b] for b in INPUTS], axis=1)\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",
        "wrapped_model = MyModel(MyPreprocessing(), model)"
      ],
      "metadata": {
        "id": "wjH9fIqkKuW6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Compile the model with the specified loss function.\n",
        "wrapped_model.compile(optimizer=tf.keras.optimizers.Adam(),\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "# Fit the model to the training data.\n",
        "wrapped_model.fit(x=parsed_dataset.batch(4), epochs=5)"
      ],
      "metadata": {
        "id": "00LogquoMVVK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pa4ex_4eKiyb"
      },
      "source": [
        "## Check model accuracy on the test set\n",
        "\n",
        "Now that we have a trained model, we can evaluate it using the test dataset.  To do that, read and prepare the test dataset in the same way as the training dataset.  Here we specify a batch size of 1 so that each example in the test set is used exactly once to compute model accuracy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tE6d7FsrMa1p",
        "cellView": "code"
      },
      "source": [
        "test_dataset = (\n",
        "  tf.data.TFRecordDataset(TEST_FILE_PATH, compression_type='GZIP')\n",
        "    .map(parse_tfrecord, num_parallel_calls=5)\n",
        "    .batch(1))\n",
        "\n",
        "wrapped_model.evaluate(test_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nhHrnv3VR0DU"
      },
      "source": [
        "# Use the trained model to classify an image from Earth Engine\n",
        "\n",
        "Now it's time to classify the image that was exported from Earth Engine.  If the exported image is large, it will be split into multiple TFRecord files in its destination folder.  There will also be a JSON sidecar file called \"the mixer\" that describes the format and georeferencing of the image.  Here we will find the image files and the mixer file, getting some info out of the mixer that will be useful during model inference."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nmTayDitZgQ5"
      },
      "source": [
        "## Find the image files and mixer file in Cloud Storage\n",
        "\n",
        "Use `gsutil` to locate the files of interest in the output Cloud Storage bucket.  Check to make sure your image export task finished before running the following."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oUv9WMpcVp8E"
      },
      "source": [
        "# Get a list of all the files in the output bucket.\n",
        "files_list = !gsutil ls 'gs://'{OUTPUT_BUCKET}\n",
        "# Get only the files generated by the image export.\n",
        "exported_files_list = [s for s in files_list if PREDICTION_INPUT_PREFIX in s]\n",
        "\n",
        "# Get the list of image files and the JSON mixer file.\n",
        "image_files_list = []\n",
        "mixer_file = None\n",
        "for f in exported_files_list:\n",
        "  if f.endswith('.tfrecord.gz'):\n",
        "    image_files_list.append(f)\n",
        "  elif f.endswith('.json'):\n",
        "    mixer_file = f\n",
        "\n",
        "# Make sure the files are in the right order.\n",
        "image_files_list.sort()\n",
        "\n",
        "pprint(image_files_list)\n",
        "print(mixer_file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RcjYG9fk53xL"
      },
      "source": [
        "## Read the mixer file\n",
        "\n",
        "The mixer contains metadata and georeferencing information for the exported patches, each of which is in a different file.  Read the mixer to get some information needed for prediction."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gn7Dr0AAd93_"
      },
      "source": [
        "# Load the contents of the mixer file to a JSON object.\n",
        "mixer_text = !gsutil cat {mixer_file}\n",
        "# Get a single string w/ newlines from the IPython.utils.text.SList\n",
        "mixer = json.loads(mixer_text.nlstr)\n",
        "pprint(mixer)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xyzyPPJwpVI"
      },
      "source": [
        "## Read the images on which to perform prediction into a dataset\n",
        "\n",
        "You can feed the list of files (`imageFilesList`) directly to the `TFRecordDataset` constructor to make a combined dataset on which to perform inference.  The input needs to be preprocessed differently than the training and testing.  Mainly, this is because the pixels are written into records as patches, we need to read the patches in as one big tensor (one patch for each band), then flatten them into lots of little tensors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tn8Kj3VfwpiJ",
        "cellView": "code"
      },
      "source": [
        "# Get relevant info from the JSON mixer file.\n",
        "patch_width = mixer['patchDimensions'][0]\n",
        "patch_height = mixer['patchDimensions'][1]\n",
        "patches = mixer['totalPatches']\n",
        "patch_dimensions_flat = [patch_width * patch_height, 1]\n",
        "\n",
        "# Note that the tensors are in the shape of a patch, one patch for each band.\n",
        "image_columns = [\n",
        "  tf.io.FixedLenFeature(shape=patch_dimensions_flat, dtype=tf.float32)\n",
        "    for k in BANDS\n",
        "]\n",
        "\n",
        "# Parsing dictionary.\n",
        "image_features_dict = dict(zip(BANDS, image_columns))\n",
        "\n",
        "# Note that you can make one dataset from many files by specifying a list.\n",
        "image_dataset = tf.data.TFRecordDataset(image_files_list, compression_type='GZIP')\n",
        "\n",
        "# Inspect a single record.\n",
        "record = iter(image_dataset).next().numpy()\n",
        "ex = tf.train.Example.FromString(record)\n",
        "print(list(ex.features.feature.keys()))\n",
        "print(len(ex.features.feature['SR_B1'].float_list.value))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Parsing function.\n",
        "def parse_image(example_proto):\n",
        "  return tf.io.parse_example(example_proto, image_features_dict)\n",
        "\n",
        "# Parse the data into tensors, one long tensor per patch.\n",
        "prediction_dataset = image_dataset.map(parse_image, num_parallel_calls=5)"
      ],
      "metadata": {
        "id": "EIOD8cjKKZp3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_2sfRemRRDkV"
      },
      "source": [
        "## Generate predictions for the image pixels\n",
        "\n",
        "To get predictions in each pixel, run the image dataset through the trained model using `model.predict()`.  Print the first prediction to see that the output is a list of the three class probabilities for each pixel.  Running all predictions might take a while."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8VGhmiP_REBP"
      },
      "source": [
        "# Run prediction in batches, with as many steps as there are patches.\n",
        "predictions = wrapped_model.predict(prediction_dataset, verbose=1)\n",
        "\n",
        "# Note that the predictions come as a numpy array.  Check the first one.\n",
        "print(predictions.shape)\n",
        "print(predictions[0])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bPU2VlPOikAy"
      },
      "source": [
        "## Write the predictions to a TFRecord file\n",
        "\n",
        "Now that there's a list of class probabilities in `predictions`, it's time to write them back into a file, optionally including a class label which is simply the index of the maximum probability.  We'll write directly from TensorFlow to a file in the output Cloud Storage bucket.\n",
        "\n",
        "Iterate over the list, compute class label and write the class and the probabilities in patches.  Specifically, we need to write the pixels into the file as patches in the same order they came out.  The records are written as serialized `tf.train.Example` protos.  This might take a while."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AkorbsEHepzJ"
      },
      "source": [
        "print('Writing to file ' + PREDICTION_OUTPUT_FILE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kATMknHc0qeR",
        "cellView": "code"
      },
      "source": [
        "def get_empty_example():\n",
        "  return tf.train.Example(\n",
        "      features=tf.train.Features(\n",
        "        feature={\n",
        "          'label': tf.train.Feature(\n",
        "              float_list=tf.train.FloatList(\n",
        "                  value=[])),\n",
        "          'bareProb': tf.train.Feature(\n",
        "              float_list=tf.train.FloatList(\n",
        "                  value=[])),\n",
        "          'vegProb': tf.train.Feature(\n",
        "              float_list=tf.train.FloatList(\n",
        "                  value=[])),\n",
        "          'waterProb': tf.train.Feature(\n",
        "              float_list=tf.train.FloatList(\n",
        "                  value=[])),\n",
        "        }\n",
        "      )\n",
        "    )\n",
        "\n",
        "# Instantiate the writer.\n",
        "writer = tf.io.TFRecordWriter(PREDICTION_OUTPUT_FILE)\n",
        "\n",
        "# Every patch-worth of predictions we'll dump an example into the output\n",
        "# file with a single feature that holds our predictions. Since our predictions\n",
        "# are already in the order of the exported data, the patches we create here\n",
        "# will also be in the right order.\n",
        "patch = 1\n",
        "pred = 0\n",
        "e = get_empty_example()\n",
        "for prediction in predictions:\n",
        "  e.features.feature['label'].float_list.value.extend([tf.argmax(prediction, -1)])\n",
        "  e.features.feature['bareProb'].float_list.value.extend([prediction[0]])\n",
        "  e.features.feature['vegProb'].float_list.value.extend([prediction[1]])\n",
        "  e.features.feature['waterProb'].float_list.value.extend([prediction[2]])\n",
        "  pred += 1\n",
        "  if (pred == patch_width * patch_height):\n",
        "    print('Done with patch ' + str(patch) + ' of ' + str(patches) + '...')\n",
        "    writer.write(e.SerializeToString())\n",
        "    patch += 1\n",
        "    pred = 0\n",
        "    e = get_empty_example()\n",
        "\n",
        "writer.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1K_1hKs0aBdA"
      },
      "source": [
        "# Upload the classifications to an Earth Engine asset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M6sNZXWOSa82"
      },
      "source": [
        "## Verify the existence of the predictions file\n",
        "\n",
        "At this stage, there should be a predictions TFRecord file sitting in the output Cloud Storage bucket.  Use the `gsutil` command to verify that the predictions image (and associated mixer JSON) exist and have non-zero size."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6ZVWDPefUCgA"
      },
      "source": [
        "!gsutil ls -l {PREDICTION_OUTPUT_FILE}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ZyCo297Clcx"
      },
      "source": [
        "## Upload the classified image to Earth Engine\n",
        "\n",
        "Upload the image to Earth Engine directly from the Cloud Storage bucket with the [`earthengine` command](https://developers.google.com/earth-engine/command_line#upload).  Provide both the image TFRecord file and the JSON file as arguments to `earthengine upload`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NXulMNl9lTDv",
        "cellView": "code"
      },
      "source": [
        "print('Uploading to ' + PREDICTION_ASSET_ID)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Of course, you also need to authenticate to the command line.\n",
        "!earthengine authenticate --auth_mode=notebook"
      ],
      "metadata": {
        "id": "Xbew3pYMWW3-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V64tcVxsO5h6"
      },
      "source": [
        "# Start the upload.\n",
        "!earthengine upload image --asset_id={PREDICTION_ASSET_ID} --pyramiding_policy=mode {PREDICTION_OUTPUT_FILE} {mixer_file}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yt4HyhUU_Bal"
      },
      "source": [
        "## Check the status of the asset ingestion\n",
        "\n",
        "You can also use the Earth Engine API to check the status of your asset upload.  It might take a while.  The upload of the image is an asset ingestion task."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_vB-gwGhl_3C",
        "cellView": "code"
      },
      "source": [
        "ee.batch.Task.list()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vvXvy9GDhM-p"
      },
      "source": [
        "## View the ingested asset\n",
        "\n",
        "Display the vector of class probabilities as an RGB image with colors corresponding to the probability of bare, vegetation, water in a pixel.  Also display the winning class using the same color palette."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kEkVxIyJiFd4"
      },
      "source": [
        "predictions_image = ee.Image(PREDICTION_ASSET_ID)\n",
        "\n",
        "prediction_vis = {\n",
        "  'bands': 'label',\n",
        "  'min': 0,\n",
        "  'max': 2,\n",
        "  'palette': ['red', 'green', 'blue']\n",
        "}\n",
        "probability_vis = {'bands': ['bareProb', 'vegProb', 'waterProb'], 'max': 0.5}\n",
        "\n",
        "prediction_map_id = predictions_image.getMapId(prediction_vis)\n",
        "probability_map_id = predictions_image.getMapId(probability_vis)\n",
        "\n",
        "map = folium.Map(location=[37.6413, -122.2582])\n",
        "folium.TileLayer(\n",
        "  tiles=prediction_map_id['tile_fetcher'].url_format,\n",
        "  attr=ATTRIBUTION,\n",
        "  overlay=True,\n",
        "  name='prediction',\n",
        ").add_to(map)\n",
        "folium.TileLayer(\n",
        "  tiles=probability_map_id['tile_fetcher'].url_format,\n",
        "  attr=ATTRIBUTION,\n",
        "  overlay=True,\n",
        "  name='probability',\n",
        ").add_to(map)\n",
        "map.add_child(folium.LayerControl())\n",
        "map"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}