guides/linked/Earth_Engine_training_patches_getPixels.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "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": [
        "<table class=\"ee-notebook-buttons\" align=\"left\"><td>\n",
        "<a target=\"_blank\"  href=\"http://colab.research.google.com/github/google/earthengine-community/blob/master/guides/linked/Earth_Engine_training_patches_getPixels.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /> Run in Google Colab</a>\n",
        "</td><td>\n",
        "<a target=\"_blank\"  href=\"https://github.com/google/earthengine-community/blob/master/guides/linked/Earth_Engine_training_patches_getPixels.ipynb\"><img width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /> View source on GitHub</a></td></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Download training patches from Earth Engine\n",
        "\n",
        "This demonstration shows how to get patches of imagery from Earth Engine assets.  Specifically, use `getPixels` calls in parallel to write a TFRecord file."
      ],
      "metadata": {
        "id": "9SV-E0p6PpGr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Imports"
      ],
      "metadata": {
        "id": "uvlyhBESPQKW"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rppQiHjZPX_y"
      },
      "outputs": [],
      "source": [
        "import concurrent\n",
        "import ee\n",
        "import google\n",
        "import io\n",
        "import json\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.animation as animation\n",
        "import multiprocessing\n",
        "import numpy as np\n",
        "import requests\n",
        "import tensorflow as tf\n",
        "\n",
        "from google.api_core import retry\n",
        "from google.colab import auth\n",
        "from google.protobuf import json_format\n",
        "from IPython.display import Image\n",
        "from matplotlib import rc\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "rc('animation', html='html5')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Authentication and initialization\n",
        "\n",
        "Use the Colab auth widget to get credentials, then use them to initialize Earth Engine.  During initialization, be sure to specify a project and Earth Engine's [high-volume endpoint](https://developers.google.com/earth-engine/cloud/highvolume), in order to make automated requests."
      ],
      "metadata": {
        "id": "pbLzoz4klKwH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# REPLACE WITH YOUR PROJECT!\n",
        "PROJECT = 'your-project'"
      ],
      "metadata": {
        "id": "HN5H25U_JBdp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "auth.authenticate_user()"
      ],
      "metadata": {
        "id": "TLmI05-wT_GD"
      },
      "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",
      "source": [
        "## Define variables"
      ],
      "metadata": {
        "id": "q7rHLQsPuwyb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# REPLACE WITH YOUR BUCKET!\n",
        "OUTPUT_FILE = 'gs://your-bucket/your-file.tfrecord.gz'\n",
        "\n",
        "# MODIS vegetation indices, 16-day.\n",
        "MOD13Q1 = ee.ImageCollection('MODIS/061/MOD13Q1').select('NDVI')\n",
        "\n",
        "# Output resolution in meters.\n",
        "SCALE = 250\n",
        "\n",
        "# Bay area.\n",
        "ROI = ee.Geometry.Rectangle(\n",
        "    [-123.05832753906247, 37.03109527141115,\n",
        "     -121.14121328124997, 38.24468432993584])\n",
        "\n",
        "# Number of samples per ROI, per year, and per TFRecord file.\n",
        "N = 64\n",
        "\n",
        "# A random sample of N locations in the ROI as a list of GeoJSON points.\n",
        "SAMPLE = ee.FeatureCollection.randomPoints(\n",
        "    region=ROI, points=N, maxError=1).aggregate_array('.geo').getInfo()\n",
        "\n",
        "# The years from which to sample every 16-day composite.\n",
        "YEARS = np.arange(2010, 2023)"
      ],
      "metadata": {
        "id": "hj_ZujvvFlGR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Image retrieval functions\n",
        "\n",
        "This section has a function to get a 1000x1000 meter patch of pixels from an asset, centered on the provided coordinates, as a numpy array.  The function can be retried automatically by using the [Retry](https://googleapis.dev/python/google-api-core/latest/retry.html) decorator.  There is also a function to serialize a structured array to a `tf.Example` proto."
      ],
      "metadata": {
        "id": "vbEM4nlUOmQn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@retry.Retry()\n",
        "def get_patch(coords, asset_id, band):\n",
        "  \"\"\"Get a patch of pixels from an asset, centered on the coords.\"\"\"\n",
        "  point = ee.Geometry.Point(coords)\n",
        "  request = {\n",
        "    'fileFormat': 'NPY',\n",
        "    'bandIds': [band],\n",
        "    'region': point.buffer(1000).bounds().getInfo(),\n",
        "    'assetId': asset_id\n",
        "  }\n",
        "  return np.load(io.BytesIO(ee.data.getPixels(request)))[band]\n",
        "\n",
        "\n",
        "def _float_feature(floats):\n",
        "  \"\"\"Returns a float_list from a float list.\"\"\"\n",
        "  return tf.train.Feature(float_list=tf.train.FloatList(value=floats))\n",
        "\n",
        "\n",
        "def array_to_example(struct_array):\n",
        "  \"\"\"\"Serialize a structured numpy array into a tf.Example proto.\"\"\"\n",
        "  struct_names = struct_array.dtype.names\n",
        "  feature = {}\n",
        "  shape = np.shape(struct_array[struct_names[0]])\n",
        "  feature['h'] = _float_feature([shape[1]])\n",
        "  feature['w'] = _float_feature([shape[2]])\n",
        "  for f in struct_names:\n",
        "    feature[f] = _float_feature(struct_array[f].flatten())\n",
        "  return tf.train.Example(\n",
        "      features = tf.train.Features(feature = feature))"
      ],
      "metadata": {
        "id": "NeKS5M-kRT4r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Get patches from the images\n",
        "\n",
        "In the variable declarations, there's a random sample in an arbitrary region of interest and a year range.  At each point in the sample, in each year, in each 16-day composite, get a patch.  The patch extraction is handled in multiple threads using a `ThreadPoolExecutor`.  Write into TFRecords where each record stores all patches for a (point, year) combination."
      ],
      "metadata": {
        "id": "H-3tJLl4WRkS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "executor = concurrent.futures.ThreadPoolExecutor(max_workers=200)\n",
        "\n",
        "writer = tf.io.TFRecordWriter(OUTPUT_FILE, 'GZIP')\n",
        "\n",
        "for point in tqdm(SAMPLE):\n",
        "  for year in tqdm(YEARS):\n",
        "    year = int(year)\n",
        "    images = MOD13Q1.filter(\n",
        "        ee.Filter.calendarRange(year, year, 'year')).getInfo()['features']\n",
        "\n",
        "    future_to_image = {\n",
        "        executor.submit(get_patch, point['coordinates'], image['id'], 'NDVI'):\n",
        "            image['id'] for image in images\n",
        "    }\n",
        "\n",
        "    arrays = ()\n",
        "    types = []\n",
        "    for future in concurrent.futures.as_completed(future_to_image):\n",
        "      image_id = future_to_image[future]\n",
        "      image_name = image_id.split('/')[-1]\n",
        "      try:\n",
        "          np_array = future.result()\n",
        "          arrays += (np_array,)\n",
        "          types.append((image_name, np.int_, np_array.shape))\n",
        "      except Exception as e:\n",
        "          print(e)\n",
        "          pass\n",
        "    array = np.array([arrays], types)\n",
        "    example_proto = array_to_example(array)\n",
        "    writer.write(example_proto.SerializeToString())\n",
        "    writer.flush()\n",
        "\n",
        "writer.close()"
      ],
      "metadata": {
        "id": "Hs_FozNIQFXI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Inspect the written files\n",
        "\n",
        "The parsing function dynamically determines the shape and keys of each record, which may vary by point and year.  Once the data are parsed, they can be displayed as an animation: one year's worth of NDVI change in a patch centered on the point."
      ],
      "metadata": {
        "id": "-ecw45dSPF-J"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "h_col = tf.io.FixedLenFeature(shape=(1), dtype=tf.float32)\n",
        "w_col = tf.io.FixedLenFeature(shape=(1), dtype=tf.float32)\n",
        "hw_dict = {'h': h_col, 'w': w_col}\n",
        "\n",
        "def parse_tfrecord(example_proto):\n",
        "  \"\"\"Parse a serialized example, dynamic determination of shape and keys.\"\"\"\n",
        "  hw = tf.io.parse_single_example(example_proto, hw_dict)\n",
        "  h = int(hw['h'].numpy())\n",
        "  w = int(hw['w'].numpy())\n",
        "\n",
        "  example = tf.train.Example()\n",
        "  example.ParseFromString(example_proto.numpy())\n",
        "  f_list = list(example.features.feature.keys())\n",
        "  f_dict = {e: tf.io.FixedLenFeature(shape=(h,w), dtype=tf.float32) for e in f_list if e not in ('h', 'w')}\n",
        "  return tf.io.parse_single_example(example_proto, f_dict)"
      ],
      "metadata": {
        "id": "wxcbEKQxaua9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = tf.data.TFRecordDataset(OUTPUT_FILE, compression_type='GZIP')\n",
        "parsed_data = [parse_tfrecord(rec) for rec in dataset]"
      ],
      "metadata": {
        "id": "v0eA8-ikbJUP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Get an animation of the data in a record\n",
        "\n",
        "See [this reference](https://matplotlib.org/stable/gallery/animation/dynamic_image.html) for details, including options to save the animation."
      ],
      "metadata": {
        "id": "w6JK6FUXMSkE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "array_dict = parsed_data[400]\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "\n",
        "# This order.\n",
        "images_names = np.sort(list(array_dict.keys()))\n",
        "first_image = images_names[0]\n",
        "\n",
        "ax.imshow(np.squeeze(array_dict[first_image]))  # show an initial one first\n",
        "ims = []\n",
        "for image in images_names[1:]:\n",
        "    im = ax.imshow(np.squeeze(array_dict[image]), animated=True)\n",
        "    ims.append([im])\n",
        "\n",
        "ani = animation.ArtistAnimation(fig, ims, interval=100, blit=True, repeat_delay=1000)\n",
        "ani"
      ],
      "metadata": {
        "id": "YoaiLbl9GCfz"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}