From 3951135893502f2ce0872d2ab790d5aea11a53ee Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:47:08 +0800 Subject: [PATCH 1/9] Add a plain-language start guide and three reader notebooks docs/start-here.md explains what TensorEE does and does not do, gives three ways to begin (the Code Editor demo, a Python quickstart, bring your own model) and defines chip, tile, halo, weight asset, export and EECU cost. README.md and docs/README.md link it first. notebooks/01_quickstart.ipynb runs the trained water U-Net on one chip at a place the reader picks, from the Code Editor module's template and the public weight tables, so it needs no checkpoint. Offline, its graph for the preset chip is byte-identical to the recorded predict_segwater graph. notebooks/02_is_my_model_supported.ipynb checks models offline: a supported U-Net, refused LinkNet and DeepLabV3+, the converter's report and a refusal, and a cell for the reader's own model. notebooks/03_bring_your_own_cnn.ipynb follows docs/own-model.md from an SMP Unet checkpoint to a region, with every billed step behind its own switch. tests/test_notebooks.py reads the notebooks with json, checks that outputs are cleared and no private values appear, and runs the code cells offline with every Earth Engine data call refused. --- README.md | 4 + docs/README.md | 8 +- docs/start-here.md | 101 ++++ notebooks/01_quickstart.ipynb | 617 ++++++++++++++++++++ notebooks/02_is_my_model_supported.ipynb | 467 +++++++++++++++ notebooks/03_bring_your_own_cnn.ipynb | 693 +++++++++++++++++++++++ tests/test_notebooks.py | 187 ++++++ 7 files changed, 2075 insertions(+), 2 deletions(-) create mode 100644 docs/start-here.md create mode 100644 notebooks/01_quickstart.ipynb create mode 100644 notebooks/02_is_my_model_supported.ipynb create mode 100644 notebooks/03_bring_your_own_cnn.ipynb create mode 100644 tests/test_notebooks.py diff --git a/README.md b/README.md index 41071da..5d99749 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,10 @@ shown alongside it uses a separately deployed model endpoint. ## Get started +- **[Start here](https://github.com/noelEOS/tensor-ee/blob/main/docs/start-here.md).** New to TensorEE? + A plain-language guide with three notebooks that run in Colab: a quickstart with the + trained U-Net, a check of whether your own model is supported, and your own CNN from + checkpoint to Earth Engine. - **[Run the Code Editor demo](https://github.com/noelEOS/tensor-ee/blob/main/CODE_EDITOR.md).** No install or checkpoint. The trained U-Net reads public weights; you need an Earth Engine account and a [registered Cloud project](https://developers.google.com/earth-engine/guides/access) diff --git a/docs/README.md b/docs/README.md index d7d0d83..61852f4 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,12 +1,16 @@ # Documentation -Start with [Getting started](getting-started.md). The guides below are the reader -documentation; the records after them hold the evidence the guides cite. +New to TensorEE? Start with [Start here](start-here.md): a plain-language guide and three +notebooks. [Getting started](getting-started.md) is the reference for installing and running +it. The guides below are the reader documentation; the records after them hold the evidence +the guides cite. ## Guides | Guide | What it covers | |---|---| +| [Start here](start-here.md) | A plain-language guide: what TensorEE does, three ways to begin, and the words you will meet | +| [Notebooks](../notebooks/) | [Quickstart](../notebooks/01_quickstart.ipynb) with the trained U-Net; [is my model supported?](../notebooks/02_is_my_model_supported.ipynb), offline; [bring your own CNN](../notebooks/03_bring_your_own_cnn.ipynb), from checkpoint to a region. Each opens in Colab | | [Getting started](getting-started.md) | Install, the offline example, both demos and a first conversion, with expected output | | [Code Editor demo](../CODE_EDITOR.md) | The trained U-Net in the Earth Engine Code Editor: no install, no checkpoint | | [Browser demo](../DEMO.md) | A local web app that runs either trained model in Earth Engine, with an offline mode | diff --git a/docs/start-here.md b/docs/start-here.md new file mode 100644 index 0000000..0e8d692 --- /dev/null +++ b/docs/start-here.md @@ -0,0 +1,101 @@ +# Start here + +This page is for anyone with an Earth Engine account and some Python who wants to try +TensorEE. You do not need to know how the library is built, and you do not need any special +tools: a browser, and for the notebooks, Google Colab or Jupyter. + +## What TensorEE does, and what it does not + +TensorEE runs a trained neural network inside Earth Engine. You keep your imagery in Earth +Engine, and the model's arithmetic (convolutions, additions, activations) is written as +ordinary Earth Engine image and array operations, so Earth Engine computes the prediction +itself. There is no model server to deploy and no imagery to download. The weights are stored +as Earth Engine table assets. It works only for the kinds of model it supports: some +convolutional networks (CNNs) such as standard U-Nets, and a few attention models through a +converter. It refuses anything else with a message saying what is missing, rather than +running an approximation. It does not train models, and it does not tell you how accurate a +model is: it checks that Earth Engine computes the same numbers as PyTorch, which is a +different question from whether the predictions are right. + +## Pick a path + +### 1. See it work, with nothing to install + +The [Code Editor demo](../CODE_EDITOR.md) runs a trained water-mapping U-Net on Sentinel-1 +radar imagery. You paste two scripts into the Earth Engine Code Editor, click **Run**, and +export one 224 × 224 pixel prediction to your own project. It needs an Earth Engine account and +a Cloud project [registered for Earth Engine](https://developers.google.com/earth-engine/guides/access). +One export is billed to your project; the first recorded run took 140 s and 64.9 batch +EECU-seconds ([record](evidence/l06d-user-check.json)). + +### 2. Use it from Python + +[Notebook 01, quickstart](../notebooks/01_quickstart.ipynb) +([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/01_quickstart.ipynb)) +does the same thing from a Python notebook: it installs TensorEE, signs in with your own Earth +Engine project, finds a Sentinel-1 scene at a place you pick, runs the same trained U-Net on +one chip there, and shows the radar image and the water map on an interactive map. It uses +only the public weights; you do not need the model's checkpoint file. + +Every notebook first runs offline: it builds the Earth Engine computation without sending it +anywhere, so you can read and run it without an account. The steps that use your Earth Engine +project are switched off until you change `RUN_LIVE = False` to `RUN_LIVE = True` and put your +project ID in place of `YOUR_EE_PROJECT`. Before each of those steps the notebook says what it +will create and what it costs. + +### 3. Bring your own model + +- [Notebook 02, is my model supported?](../notebooks/02_is_my_model_supported.ipynb) + ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)) + needs no Earth Engine account at all. It loads a PyTorch model, checks it, and explains in + plain words what is supported, what is refused and why, and what you can do next. It shows a + model that passes and models that are refused, and has a cell where you put your own. +- [Notebook 03, bring your own CNN](../notebooks/03_bring_your_own_cnn.ipynb) + ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)) + follows a trained segmentation-models-pytorch U-Net from its checkpoint file to Earth + Engine: check it, plan how to cut a region into tiles, upload its weights once, predict one + tile, compare that tile with PyTorch, and export a region. Each step that is billed to your + project is switched off until you turn it on. + +When you want the details, the [Getting started](getting-started.md) page and the +[guides](README.md) are the reference. + +## Words you will meet + +- **Chip.** One fixed-size square of pixels that a model takes as input, for example + 224 × 224 pixels of Sentinel-1 at 10 m. The trained U-Net was trained on chips of that size. +- **Tile.** One piece of a larger region. A region too big for one run is cut into tiles that + Earth Engine computes one by one, and the results are put back together. +- **Halo.** The extra border of input pixels around a tile that a model needs to compute the + tile's edge correctly, because each output pixel depends on input pixels around it. The halo + is computed but thrown away, so it is a cost: for a standard U-Net with a ResNet-34 encoder + it is 544 pixels on each side, so a 1,024-pixel tile reads a 2,112-pixel square. +- **Weight asset.** An Earth Engine table asset that holds a model's numbers (its weights). + You upload them once, and every prediction reads them from there. The trained U-Net's weight + assets are public, so you can use them without uploading anything. +- **Export.** An Earth Engine batch task (`Export.image.toAsset` or `toDrive`) that computes an + image and saves it as an asset or a file. Whole models run as exports, because one + interactive request with a whole model exceeded Earth Engine's memory limit. An export runs + in the background: you start it, wait for **Completed** in the Tasks list, then look at the + result. +- **EECU cost.** Earth Engine measures computation in EECU-seconds (Earth Engine Compute + Units), and bills commercial projects for them; see Google's + [Earth Engine pricing](https://cloud.google.com/earth-engine/pricing). Batch exports and + interactive requests are counted separately. The notebooks give the recorded cost where one exists, and say so when a cost has not + been measured. + +## What has been checked + +TensorEE keeps four kinds of claim apart, and the notebooks do too: + +- **Graph identity**: the Earth Engine computation is exactly, byte for byte, one that was + checked before. +- **Offline agreement**: each building block, computed in NumPy, matches PyTorch. +- **Live agreement**: Earth Engine's result matches PyTorch on the same input, within an + absolute difference of 1e-4 in the logits. For the trained U-Net the recorded single-chip + run differed by at most 1.35e-5 ([validation](validation.md)). +- **Accuracy**: how often the predictions are right. This is not measured by TensorEE. You + measure it on your own labelled data. + +If a notebook gets a result you did not expect, or refuses a model you think should work, +[open an issue](https://github.com/noelEOS/tensor-ee/issues) with the message it printed. diff --git a/notebooks/01_quickstart.ipynb b/notebooks/01_quickstart.ipynb new file mode 100644 index 0000000..330418e --- /dev/null +++ b/notebooks/01_quickstart.ipynb @@ -0,0 +1,617 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2ac26326", + "metadata": {}, + "source": [ + "# TensorEE quickstart: a trained water U-Net on one Sentinel-1 chip\n", + "\n", + "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/01_quickstart.ipynb)\n", + "\n", + "This notebook runs a trained neural network inside Google Earth Engine. The model is a U-Net\n", + "with a ResNet-18 encoder that maps open water in Sentinel-1 radar imagery. It takes one chip of\n", + "224 × 224 pixels at 10 m (2.24 km on a side) and returns, for every pixel, the probability\n", + "that it is water and a water mask.\n", + "\n", + "Earth Engine computes the whole prediction from your Sentinel-1 scene. Nothing is downloaded\n", + "and no model server is involved: the network's arithmetic is a large Earth Engine computation,\n", + "and its weights are public Earth Engine table assets. You do not need the model's checkpoint\n", + "file.\n", + "\n", + "**What you need.** Nothing, to read and run the notebook offline. To run the model, an Earth\n", + "Engine account and a Cloud project\n", + "[registered for Earth Engine](https://developers.google.com/earth-engine/guides/access) where\n", + "you can create an image asset.\n", + "\n", + "**What it costs.** Offline, nothing. With `RUN_LIVE = True` the notebook makes a few small\n", + "requests (to place the chip and find a scene) and **one batch export of one chip**, billed to\n", + "your project. The first recorded export of this chip took 140 s and 64.9 batch EECU-seconds.\n", + "The export creates one image asset in your project (about 1 MB). Each step says what it does\n", + "before it does it." + ] + }, + { + "cell_type": "markdown", + "id": "7361375f", + "metadata": {}, + "source": [ + "## 1. Install\n", + "\n", + "The trained U-Net's Earth Engine computation is stored in the repository, in the file the\n", + "[Code Editor demo](https://github.com/noelEOS/tensor-ee/blob/main/CODE_EDITOR.md) uses. That\n", + "file is not part of the package on PyPI, so this cell copies the repository and installs\n", + "TensorEE from it. It also installs `folium` for the map. In Colab it takes a minute or two.\n", + "\n", + "If you opened this notebook from your own copy of the repository, you can skip this cell." + ] + }, + { + "cell_type": "code", + "id": "33d242a2", + "metadata": {}, + "source": [ + "# The Code Editor module (demos/code_editor/segwater.js) is in the repository, not on PyPI,\n", + "# so this notebook installs TensorEE from a copy of the repository.\n", + "!git clone --depth 1 https://github.com/noelEOS/tensor-ee.git\n", + "%pip install --quiet ./tensor-ee folium" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "9b2cdbb8", + "metadata": {}, + "source": [ + "import sys\n", + "from pathlib import Path\n", + "\n", + "\n", + "def find_repository():\n", + " \"\"\"The copy of the repository: ./tensor-ee in Colab, or the folder this notebook is in.\"\"\"\n", + " for folder in (Path(\"tensor-ee\"), Path(\".\"), Path(\"..\")):\n", + " if (folder / \"demos\" / \"code_editor\" / \"segwater.js\").is_file():\n", + " return folder.resolve()\n", + " raise FileNotFoundError(\"Run the install cell first: it copies the repository to ./tensor-ee\")\n", + "\n", + "\n", + "REPOSITORY = find_repository()\n", + "sys.path.insert(0, str(REPOSITORY))\n", + "print(\"Using the repository in\", REPOSITORY)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "376fa9ae", + "metadata": {}, + "source": [ + "## 2. Settings\n", + "\n", + "This is the only cell you need to edit.\n", + "\n", + "- `PROJECT`: your Cloud project ID, the one you use in the Code Editor.\n", + "- `RUN_LIVE`: leave it `False` to run everything offline first. Set it to `True` to sign in\n", + " and run the steps that use your project.\n", + "- `LONGITUDE`, `LATITUDE`: the centre of the chip, in degrees. The default is Sun Moon Lake,\n", + " Taiwan, where the model's Earth Engine result was checked against PyTorch.\n", + "- `START_DATE`, `END_DATE`: the period in which to look for a Sentinel-1 scene.\n", + "- `ASSET_NAME`: the name of the new image asset the export creates in your project. Use a new\n", + " name for each run: an export never overwrites an asset.\n", + "\n", + "After changing `RUN_LIVE`, restart the runtime (in Colab: **Runtime → Restart session**) and\n", + "run the notebook again from the top: the offline mode stays installed until the runtime\n", + "restarts." + ] + }, + { + "cell_type": "code", + "id": "c035f536", + "metadata": {}, + "source": [ + "PROJECT = \"YOUR_EE_PROJECT\" # your Cloud project registered for Earth Engine\n", + "RUN_LIVE = False # True: sign in and run the steps that use your project\n", + "\n", + "LONGITUDE, LATITUDE = 120.9150, 23.8620 # Sun Moon Lake, Taiwan\n", + "START_DATE, END_DATE = \"2024-01-01\", \"2024-01-08\" # look for Sentinel-1 scenes in this period\n", + "ASSET_NAME = \"tensor_ee_quickstart_chip\" # a new image asset in your project" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "5b6ecfc4", + "metadata": {}, + "source": [ + "PLACEHOLDER = \"YOUR_EE_PROJECT\"\n", + "\n", + "\n", + "def live_step(what):\n", + " \"\"\"True when a step that uses Earth Engine should run; it refuses a placeholder project.\"\"\"\n", + " if not RUN_LIVE:\n", + " print(f\"Skipped: {what}. RUN_LIVE is False, so nothing was sent to Earth Engine.\")\n", + " return False\n", + " if not PROJECT or PROJECT == PLACEHOLDER:\n", + " raise ValueError(\"Set PROJECT to your own Cloud project ID in the settings cell \"\n", + " \"before running with RUN_LIVE = True.\")\n", + " return True" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "d32406d8", + "metadata": {}, + "source": [ + "## 3. Sign in, or work offline\n", + "\n", + "With `RUN_LIVE = True` the next cell signs you in. The first time, `ee.Authenticate()` shows\n", + "a link or a prompt; follow it with the Google account you use for Earth Engine. Every request\n", + "after that is made with your account and billed to `PROJECT`.\n", + "\n", + "With `RUN_LIVE = False`, the cell after it installs the Earth Engine client's offline\n", + "function signatures (`ee.apitestcase.ApiTestCase`, a test helper that ships with the client).\n", + "They let the notebook build the same Earth Engine computation without credentials and without\n", + "sending anything." + ] + }, + { + "cell_type": "code", + "id": "ee640c02", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "import ee\n", + "\n", + "if live_step(\"sign in to Earth Engine\"):\n", + " ee.Authenticate()\n", + " ee.Initialize(project=PROJECT)\n", + " print(\"Signed in. Requests are billed to\", PROJECT)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "d1822a6b", + "metadata": {}, + "source": [ + "import ee\n", + "\n", + "if not RUN_LIVE:\n", + " from ee.apitestcase import ApiTestCase\n", + "\n", + " offline = ApiTestCase()\n", + " offline.setUp() # offline signatures: no credentials, no network\n", + " print(\"Offline: the notebook builds Earth Engine computations but sends none.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "6b906da1", + "metadata": {}, + "source": [ + "## 4. Load the trained model\n", + "\n", + "The model comes as one file, `demos/code_editor/segwater.js`. TensorEE built it from the\n", + "trained checkpoint: it holds the complete Earth Engine computation for one chip, with blanks\n", + "for the place and the scene, and it reads the network's weights from 29 public Earth Engine\n", + "table assets. That is why you need neither the checkpoint nor an upload. The same file runs in\n", + "the Code Editor demo.\n", + "\n", + "The file stores the computation compressed. Reading it checks that it is complete and\n", + "undamaged." + ] + }, + { + "cell_type": "code", + "id": "e68ae525", + "metadata": {}, + "source": [ + "from demos.code_editor import build as segwater_module\n", + "\n", + "TEMPLATE, METADATA = segwater_module.parse_module(\n", + " (REPOSITORY / \"demos\" / \"code_editor\" / \"segwater.js\").read_text())\n", + "\n", + "print(\"Model computation:\", f\"{METADATA['templateBytes']:,}\", \"characters, checked intact\")\n", + "print(\"Weight tables it reads:\", len(METADATA[\"weightAssets\"]))\n", + "print(\"Built from the checkpoint with SHA-256\", METADATA[\"checkpointSha256\"][:16] + \"...\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "52b594cb", + "metadata": {}, + "source": [ + "## 5. Choose the place\n", + "\n", + "The chip is a 224 × 224 grid of 10 m pixels in the local UTM zone, centred on your point.\n", + "The model was trained on chips of exactly this size, so the notebook keeps it.\n", + "\n", + "Offline, the notebook uses the preset chip at Sun Moon Lake. With `RUN_LIVE = True` the live\n", + "cell asks Earth Engine to convert your longitude and latitude to UTM metres: **one small\n", + "request**." + ] + }, + { + "cell_type": "code", + "id": "be96d487", + "metadata": {}, + "source": [ + "from tensor_ee import Grid\n", + "\n", + "CHIP, PIXEL = 224, 10 # pixels per side, metres per pixel\n", + "PRESET = { # the verified Sun Moon Lake chip of the Code Editor demo\n", + " \"crs\": \"EPSG:32651\", \"x\": 286030, \"y\": 2640860,\n", + " \"scene\": \"COPERNICUS/S1_GRD/S1A_IW_GRDH_1SDV_20240103T100134_20240103T100159_051941_06469E_B0BD\",\n", + "}\n", + "\n", + "\n", + "def utm_zone(longitude, latitude):\n", + " \"\"\"The EPSG code of the WGS 84 UTM zone containing a point.\"\"\"\n", + " zone = min(60, max(1, int((longitude + 180) // 6) + 1))\n", + " return f\"EPSG:{326 if latitude >= 0 else 327}{zone:02d}\"\n", + "\n", + "\n", + "place = {\"crs\": PRESET[\"crs\"], \"x\": PRESET[\"x\"], \"y\": PRESET[\"y\"]} # offline: the preset" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "33726682", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "if live_step(\"place the chip at your point (1 request)\"):\n", + " crs = utm_zone(LONGITUDE, LATITUDE)\n", + " x, y = ee.Geometry.Point([LONGITUDE, LATITUDE]).transform(crs, 0.01).coordinates().getInfo()\n", + " half = CHIP * PIXEL // 2\n", + " # The chip's upper-left corner, on whole 10 m steps.\n", + " place = {\"crs\": crs, \"x\": round((x - half) / PIXEL) * PIXEL,\n", + " \"y\": round((y + half) / PIXEL) * PIXEL}" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "61805fcd", + "metadata": {}, + "source": [ + "grid = Grid(place[\"crs\"], (PIXEL, 0, place[\"x\"], 0, -PIXEL, place[\"y\"]), CHIP, CHIP)\n", + "print(\"Chip grid:\", grid.crs, \"upper-left corner\", (place[\"x\"], place[\"y\"]),\n", + " f\"{CHIP} x {CHIP} pixels of {PIXEL} m\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "3a027ecd", + "metadata": {}, + "source": [ + "## 6. Find a Sentinel-1 scene\n", + "\n", + "The model reads the VV and VH bands of a Sentinel-1 GRD scene (interferometric wide swath, 10 m)\n", + "in decibels, as Earth Engine stores them. The scene must cover the whole chip.\n", + "\n", + "With `RUN_LIVE = True` the live cell lists up to ten scenes in your period and, for each, the\n", + "fraction of the chip it covers, and picks the first that covers it all: **one request**, which\n", + "reads only the chip's area. Offline, the notebook uses the preset scene of 3 January 2024." + ] + }, + { + "cell_type": "code", + "id": "05f143cc", + "metadata": {}, + "source": [ + "scene_id = PRESET[\"scene\"] # offline: the preset scene" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "4d94e7c5", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "if live_step(\"find a Sentinel-1 scene (1 request)\"):\n", + " area = grid.geometry()\n", + " collection = (ee.ImageCollection(\"COPERNICUS/S1_GRD\")\n", + " .filterBounds(area).filterDate(START_DATE, END_DATE)\n", + " .filter(ee.Filter.eq(\"instrumentMode\", \"IW\"))\n", + " .filter(ee.Filter.listContains(\"transmitterReceiverPolarisation\", \"VV\"))\n", + " .filter(ee.Filter.listContains(\"transmitterReceiverPolarisation\", \"VH\"))\n", + " .filter(ee.Filter.eq(\"resolution_meters\", 10))\n", + " .sort(\"system:time_start\"))\n", + "\n", + " def describe(image):\n", + " image = ee.Image(image)\n", + " valid = (image.select([\"VV\", \"VH\"]).mask().reduce(ee.Reducer.min()).gt(0)\n", + " .unmask(0, False).rename(\"valid\"))\n", + " coverage = valid.reduceRegion(reducer=ee.Reducer.mean(), geometry=area, crs=grid.crs,\n", + " crsTransform=list(grid.transform),\n", + " maxPixels=CHIP * CHIP, tileScale=2).get(\"valid\")\n", + " return ee.Dictionary({\n", + " \"id\": ee.String(\"COPERNICUS/S1_GRD/\").cat(image.get(\"system:index\")),\n", + " \"time\": image.date().format(\"YYYY-MM-dd HH:mm\"),\n", + " \"pass\": image.get(\"orbitProperties_pass\"), \"coverage\": coverage})\n", + "\n", + " scenes = collection.toList(10).map(describe).getInfo()\n", + " for scene in scenes:\n", + " print(f\"{scene['time']} {scene['pass']:<10} covers {100 * (scene['coverage'] or 0):6.2f}% {scene['id']}\")\n", + " complete = [scene for scene in scenes if (scene[\"coverage\"] or 0) >= 0.999]\n", + " if not complete:\n", + " raise RuntimeError(\"No scene in this period covers the whole chip. Widen START_DATE to \"\n", + " \"END_DATE, or move the point away from the edge of the swath.\")\n", + " scene_id = complete[0][\"id\"]\n", + " print(\"Using\", scene_id)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "ef2053c6", + "metadata": {}, + "source": [ + "## 7. Build the prediction\n", + "\n", + "This cell fills the model's blanks with your grid and scene and turns the text back into an\n", + "Earth Engine image. Nothing is computed yet: like any Earth Engine image, it is a recipe that\n", + "Earth Engine evaluates only when you export or display it.\n", + "\n", + "The recipe converts the radar values from decibels the way the model was trained, runs the\n", + "U-Net, and returns five bands:\n", + "\n", + "| Band | Meaning |\n", + "|---|---|\n", + "| `logit_0`, `logit_1` | the network's raw scores for \"not water\" and \"water\" |\n", + "| `water_probability` | the probability of water, from 0 to 1 |\n", + "| `water_mask` | 1 where the probability is at least 0.5, else 0 |\n", + "| `coverage` | 1 where both radar bands and the result are valid |\n", + "\n", + "For the preset chip, the notebook also checks the computation against the record: it should\n", + "be identical, byte for byte, to the one TensorEE's own Python recipe builds with the trained\n", + "checkpoint, which the repository records." + ] + }, + { + "cell_type": "code", + "id": "6a971198", + "metadata": {}, + "source": [ + "import json\n", + "\n", + "sys.setrecursionlimit(max(sys.getrecursionlimit(), 20000)) # the U-Net is a deep graph\n", + "OUTPUT_BANDS = list(segwater_module.OUTPUT_BANDS)\n", + "\n", + "text = segwater_module.substitute(TEMPLATE, grid.crs, scene_id,\n", + " (place[\"x\"], place[\"y\"]), (place[\"x\"], place[\"y\"]))\n", + "chip = ee.Image(ee.deserializer.fromJSON(text))\n", + "# Put the chip on its grid, as the library's recipe does for every chip it computes.\n", + "prediction = (ee.ImageCollection.fromImages([chip]).mosaic().select(OUTPUT_BANDS)\n", + " .setDefaultProjection(grid.projection()).clip(grid.geometry())\n", + " .updateMask(grid.mask()).toFloat())\n", + "\n", + "serialized = prediction.serialize()\n", + "graph_sha256 = segwater_module.sha256(serialized)\n", + "print(\"Bands:\", OUTPUT_BANDS)\n", + "print(f\"The computation is {len(serialized):,} characters of Earth Engine expression.\")\n", + "\n", + "record = json.loads((REPOSITORY / \"docs\" / \"evidence\" / \"s03-template-identity.json\").read_text())\n", + "preset = next(aoi for aoi in record[\"aois\"] if aoi[\"name\"] == \"segwater_sun_moon_lake_1x1\")\n", + "if (scene_id, grid.crs, place[\"x\"], place[\"y\"]) != (PRESET[\"scene\"], PRESET[\"crs\"], PRESET[\"x\"], PRESET[\"y\"]):\n", + " print(\"Your own place: there is no recorded computation to compare it with.\")\n", + "elif ee.__version__ != METADATA[\"earthengineApiVersion\"]:\n", + " print(f\"Not compared: your Earth Engine client is {ee.__version__}; the record was made with \"\n", + " f\"{METADATA['earthengineApiVersion']}, and client versions write expressions differently.\")\n", + "elif graph_sha256 == preset[\"library_sha256\"]:\n", + " print(\"Identical, byte for byte, to the recorded computation of the library's recipe.\")\n", + "else:\n", + " print(\"Different from the recorded computation: please report this as an issue.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "21a05245", + "metadata": {}, + "source": [ + "## 8. Export the prediction (billed to your project)\n", + "\n", + "**What this creates.** One batch export task, `Export.image.toAsset`, and when it completes,\n", + "one new image asset, `projects/PROJECT/assets/ASSET_NAME`: five float32 bands over\n", + "224 × 224 pixels, about 1 MB.\n", + "\n", + "**What it costs.** The first recorded export of the Sun Moon Lake chip took 140 s and\n", + "64.9 batch EECU-seconds. No cost has been recorded for other places. How long an export\n", + "waits and runs varies with Earth Engine's load.\n", + "\n", + "The model runs as an export, not as map tiles: when a whole trained U-Net was sent as one\n", + "interactive request, it exceeded Earth Engine's memory limit. The cell after this one waits for the task." + ] + }, + { + "cell_type": "code", + "id": "ce479c84", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "asset_id = f\"projects/{PROJECT}/assets/{ASSET_NAME}\"\n", + "\n", + "if live_step(\"export the prediction (1 batch task, 1 new asset)\"):\n", + " payload = prediction.set({\"source_scene\": scene_id,\n", + " \"checkpoint_sha256\": METADATA[\"checkpointSha256\"],\n", + " \"template_sha256\": METADATA[\"templateSha256\"]})\n", + " task = ee.batch.Export.image.toAsset(\n", + " image=payload, description=\"tensor_ee_quickstart\", assetId=asset_id,\n", + " region=grid.geometry(), crs=grid.crs, crsTransform=list(grid.transform),\n", + " maxPixels=CHIP * CHIP * len(OUTPUT_BANDS), pyramidingPolicy={\".default\": \"sample\"})\n", + " task.start()\n", + " print(\"Started task\", task.id, \"writing\", asset_id)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "2d79e129", + "metadata": {}, + "source": [ + "Waiting checks the task's status every 20 seconds. You can also follow it in the **Tasks** tab\n", + "of the Code Editor. If you stop waiting, the export carries on; run this cell again later." + ] + }, + { + "cell_type": "code", + "id": "7e651659", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "import time\n", + "\n", + "if live_step(\"wait for the export\"):\n", + " deadline = time.time() + 30 * 60\n", + " while True:\n", + " state = task.status()[\"state\"]\n", + " if state in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", + " break\n", + " if time.time() > deadline:\n", + " raise TimeoutError(\"Still running after 30 minutes; run this cell again later.\")\n", + " time.sleep(20)\n", + " print(\"Export\", state)\n", + " if state != \"COMPLETED\":\n", + " raise RuntimeError(task.status().get(\"error_message\", \"the export did not complete\"))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "92dbf3c4", + "metadata": {}, + "source": [ + "## 9. See the radar image and the water map\n", + "\n", + "The map shows the Sentinel-1 VV band and the completed asset's water probability and mask.\n", + "The layers are Earth Engine map tiles: displaying them reads the saved asset and the scene,\n", + "so panning does not run the network again. Each tile you view is a small interactive request.\n", + "Use the layer control at the top right to switch layers." + ] + }, + { + "cell_type": "code", + "id": "f67c0038", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "if live_step(\"show the result on a map\"):\n", + " import folium\n", + " from IPython.display import display\n", + "\n", + " def tile_layer(image, visualization, name, show=True):\n", + " url = ee.Image(image).getMapId(visualization)[\"tile_fetcher\"].url_format\n", + " return folium.raster_layers.TileLayer(\n", + " tiles=url, name=name, overlay=True, show=show,\n", + " attr=\"Google Earth Engine; contains modified Copernicus Sentinel data\")\n", + "\n", + " result = ee.Image(asset_id)\n", + " outline = grid.geometry().transform(\"EPSG:4326\", 0.01).getInfo()\n", + " longitudes, latitudes = zip(*outline[\"coordinates\"][0])\n", + " water_map = folium.Map(location=[sum(latitudes) / len(latitudes), sum(longitudes) / len(longitudes)],\n", + " zoom_start=14)\n", + " tile_layer(ee.Image(scene_id).select(\"VV\"), {\"min\": -25, \"max\": 0},\n", + " \"Sentinel-1 VV (dB)\").add_to(water_map)\n", + " tile_layer(result.select(\"water_probability\"),\n", + " {\"min\": 0, \"max\": 1, \"palette\": [\"ffffff\", \"08306b\"]},\n", + " \"Water probability\").add_to(water_map)\n", + " tile_layer(result.select(\"water_mask\").selfMask(), {\"palette\": [\"00bfff\"]},\n", + " \"Water mask\", show=False).add_to(water_map)\n", + " folium.GeoJson(outline, name=\"Chip\", style_function=lambda _: {\n", + " \"fill\": False, \"color\": \"#ff7f00\", \"weight\": 2}).add_to(water_map)\n", + " folium.LayerControl().add_to(water_map)\n", + " display(water_map)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "582437dd", + "metadata": {}, + "source": [ + "## What this shows, and what it does not\n", + "\n", + "- **The computation.** For the preset chip, the notebook's computation is identical to the\n", + " library's own recipe built with the trained checkpoint (step 7). For another place it is the\n", + " same computation with a different grid and scene.\n", + "- **Earth Engine against PyTorch.** For the preset chip, an earlier export was compared pixel by\n", + " pixel with the trained model in PyTorch: the largest logit difference was 1.35e-5, the largest\n", + " probability difference 6.4e-6, and the water masks agreed on all 50,176 pixels\n", + " ([validation](https://github.com/noelEOS/tensor-ee/blob/main/docs/validation.md)). Your\n", + " export at another place has not been compared; that needs the checkpoint, which is not yet\n", + " public.\n", + "- **Accuracy.** The water map is a prediction, not ground truth, and this notebook does not\n", + " measure how often it is right. The model was trained on Sentinel-1 chips with water labels\n", + " derived from Sentinel-2\n", + " ([model notes](https://github.com/noelEOS/tensor-ee/blob/main/docs/model.md)).\n", + "- **Size.** One chip of 224 × 224 pixels per export. The Code Editor demo runs up to four chips\n", + " side by side; each chip is computed on its own, so seams can appear at chip borders.\n", + "\n", + "**Cleaning up.** The asset stays in your project until you delete it, in the Code Editor's\n", + "**Assets** tab.\n", + "\n", + "**Next.** [Notebook 02](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)\n", + "checks whether your own model can run this way, and\n", + "[notebook 03](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", + "takes a trained U-Net of your own to Earth Engine." + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/02_is_my_model_supported.ipynb b/notebooks/02_is_my_model_supported.ipynb new file mode 100644 index 0000000..6ccaadf --- /dev/null +++ b/notebooks/02_is_my_model_supported.ipynb @@ -0,0 +1,467 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "5df6cfd8", + "metadata": {}, + "source": [ + "# Is my model supported?\n", + "\n", + "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)\n", + "\n", + "TensorEE can run a model in Earth Engine only if every operation in it has a TensorEE building\n", + "block, at settings the library can compute exactly. This notebook checks that for a PyTorch\n", + "model, and explains the answer in plain words: what is supported, what is refused and why,\n", + "and what you can do next.\n", + "\n", + "**It runs entirely on your computer or in Colab.** It needs no Earth Engine account, makes no\n", + "Earth Engine request and downloads no weights: every example model is built with random\n", + "weights, which is all a support check needs. The check looks at the model's structure, never\n", + "at its predictions.\n", + "\n", + "The notebook works through four examples, then has a cell for your own model:\n", + "\n", + "1. a U-Net that is supported;\n", + "2. two models that are refused, and what the refusal says;\n", + "3. an attention model, which goes through a different path, the converter;\n", + "4. your model." + ] + }, + { + "cell_type": "markdown", + "id": "a8601d7f", + "metadata": {}, + "source": [ + "## 1. Install\n", + "\n", + "The `swin` extra brings segmentation-models-pytorch 0.5.0 and timm 1.0.30, the versions the\n", + "library was checked with. PyTorch is already installed in Colab." + ] + }, + { + "cell_type": "code", + "id": "5458608d", + "metadata": {}, + "source": [ + "%pip install --quiet \"tensor-ee[swin]\"\n", + "# Grouped and depthwise convolutions (ResNeXt and MobileNet encoders, for example) became\n", + "# supported after the 0.1.0 release. To check such a model, install from the repository instead:\n", + "# %pip install --quiet \"tensor-ee[swin] @ git+https://github.com/noelEOS/tensor-ee\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "c424214b", + "metadata": {}, + "source": [ + "## 2. Settings\n", + "\n", + "This notebook has no live step, so it has no project to set: nothing here contacts Earth\n", + "Engine. The only setting is the size of the chip the checks use." + ] + }, + { + "cell_type": "code", + "id": "295d1ada", + "metadata": {}, + "source": [ + "RUN_LIVE = False # this notebook never contacts Earth Engine\n", + "INPUT_SHAPE = (1, 2, 64, 64) # one chip: (1, bands, height, width)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "cf1e5641", + "metadata": {}, + "source": [ + "from collections import Counter\n", + "\n", + "import torch\n", + "import segmentation_models_pytorch as smp\n", + "\n", + "from tensor_ee import check_compatibility\n", + "from tensor_ee.smp import adapt_smp_model\n", + "\n", + "torch.set_num_threads(2)\n", + "torch.manual_seed(0) # random weights, the same every run\n", + "\n", + "\n", + "def explain(report):\n", + " \"\"\"Say in plain words what a compatibility report found.\"\"\"\n", + " facts = report.to_dict()\n", + " if facts[\"status\"] == \"supported\":\n", + " storage = facts[\"weight_storage\"]\n", + " print(\"SUPPORTED: every operation in this model has a TensorEE building block.\")\n", + " print(f\" {facts['parameter_count']:,} parameters. After BatchNorm is folded into the \"\n", + " f\"convolutions, {storage['coefficient_count']:,} numbers would be uploaded as weight assets.\")\n", + " print(f\" Output for one chip: {facts['output_shape']} (batch, classes, height, width).\")\n", + " return\n", + " print(f\"NOT SUPPORTED (status: {facts['status']}).\")\n", + " if not facts[\"evaluation\"][\"model_eval\"]:\n", + " print(\" The model is in training mode: call model.eval() first.\")\n", + " if not facts[\"evaluation\"][\"float32_parameters_and_buffers\"]:\n", + " print(\" Some weights are not float32: TensorEE computes in float32 at most.\")\n", + " if facts[\"trace\"][\"status\"] != \"success\":\n", + " print(\" The model's code could not be traced:\", facts[\"trace\"][\"error\"])\n", + " reasons = Counter((blocker[\"code\"], blocker[\"message\"]) for blocker in facts[\"blockers\"])\n", + " for (code, message), count in reasons.most_common(8):\n", + " print(f\" {count} x {message} [{code}]\")\n", + " if len(reasons) > 8:\n", + " print(f\" ... and {len(reasons) - 8} more kinds of refusal (report.to_dict()['blockers'] lists them all)\")\n", + " if facts[\"unknowns\"]:\n", + " print(f\" {len(facts['unknowns'])} later operations were not checked, because they depend \"\n", + " \"on the refused ones.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "00d29019", + "metadata": {}, + "source": [ + "## 3. A supported model: a U-Net with a ResNet-18 encoder\n", + "\n", + "This is the architecture of the trained water model in notebook 01, built here with random\n", + "weights (`encoder_weights=None` means nothing is downloaded).\n", + "\n", + "Two calls do the work:\n", + "\n", + "- `adapt_smp_model` prepares a segmentation-models-pytorch model for TensorEE at a fixed chip\n", + " size. It also runs the original model and the prepared one on one test input and checks that\n", + " they give exactly the same numbers.\n", + "- `check_compatibility` walks through every operation of the prepared model and says whether\n", + " TensorEE can compute it." + ] + }, + { + "cell_type": "code", + "id": "63b49161", + "metadata": {}, + "source": [ + "unet = smp.Unet(\"resnet18\", encoder_weights=None, in_channels=2, classes=2).eval()\n", + "prepared = adapt_smp_model(unet, input_shape=INPUT_SHAPE)\n", + "print(\"Prepared model matches the original on a test input:\",\n", + " prepared.tensor_ee_smp_metadata[\"torch_equal\"])\n", + "explain(check_compatibility(prepared, INPUT_SHAPE))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "443e3ab9", + "metadata": {}, + "source": [ + "\"Supported\" means the library can express this model, operation by operation, as Earth Engine\n", + "computations. It is not yet a claim about the numbers. Separately, the library has evidence\n", + "for parts of the path: the building blocks agree with PyTorch at the settings they were\n", + "validated at, and one trained U-Net with a ResNet-18 encoder, prepared this way, ran in Earth\n", + "Engine on a 64 × 64 test input and matched PyTorch to within 1.2e-5 in the logits. Other\n", + "encoders and architectures have not been run in Earth Engine as whole models. And a supported\n", + "model is not necessarily an accurate one: that depends on its training.\n", + "\n", + "Encoders checked offline in the same way: ResNet-18, 34 and 50 and VGG-11 (with BatchNorm), in\n", + "`Unet` and `UnetPlusPlus`." + ] + }, + { + "cell_type": "markdown", + "id": "0692e745", + "metadata": {}, + "source": [ + "## 4. Refused models\n", + "\n", + "### A LinkNet: transposed convolutions\n", + "\n", + "LinkNet's decoder upsamples with transposed convolutions (`ConvTranspose2d`). TensorEE has no\n", + "building block for them in a CNN, so it refuses the model and names each one." + ] + }, + { + "cell_type": "code", + "id": "944ada38", + "metadata": {}, + "source": [ + "linknet = smp.Linknet(\"resnet18\", encoder_weights=None, in_channels=2, classes=2).eval()\n", + "explain(check_compatibility(adapt_smp_model(linknet, input_shape=INPUT_SHAPE), INPUT_SHAPE))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "d07d1603", + "metadata": {}, + "source": [ + "### DeepLabV3+: dilated convolutions, and more\n", + "\n", + "DeepLabV3+ keeps its last encoder stage at a higher resolution by spacing out the taps of its\n", + "convolutions (dilation). The CNN path supports convolutions with stride 1 or 2, no dilation,\n", + "and \"same\" zero padding, so this model is refused. The report also names its other missing\n", + "operations: the adaptive average pooling and the bilinear resize of its pyramid pooling module." + ] + }, + { + "cell_type": "code", + "id": "f1d487d8", + "metadata": {}, + "source": [ + "deeplab = smp.DeepLabV3Plus(\"resnet18\", encoder_weights=None, in_channels=2, classes=2).eval()\n", + "explain(check_compatibility(adapt_smp_model(deeplab, input_shape=INPUT_SHAPE), INPUT_SHAPE))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "f73d6c10", + "metadata": {}, + "source": [ + "A refusal is a firm answer: TensorEE never replaces an operation it cannot compute exactly with\n", + "an approximation. The message names the first operations that block the model; the\n", + "operations after them are not checked. Your options:\n", + "\n", + "- **Change the architecture**, for example a U-Net decoder (nearest-neighbour upsampling and\n", + " convolutions) in place of transposed convolutions, and retrain.\n", + "- **Ask for the operation.** [Open an issue](https://github.com/noelEOS/tensor-ee/issues)\n", + " with the printed message and your model's architecture. Adding a building block is a\n", + " library change with its own tests and evidence;\n", + " [CONTRIBUTING](https://github.com/noelEOS/tensor-ee/blob/main/CONTRIBUTING.md) describes it." + ] + }, + { + "cell_type": "markdown", + "id": "b9d42f92", + "metadata": {}, + "source": [ + "## 5. Attention models: the converter\n", + "\n", + "Transformer-style models (SegFormer, Swin) do not go through `adapt_smp_model`. It refuses them\n", + "and points to the converter, `tensor_ee.converter.convert`, which recognizes the modules of the\n", + "models it was built for: SegFormer-B0 and a UPerNet with a Swin-T encoder, both on\n", + "224 × 224 chips." + ] + }, + { + "cell_type": "code", + "id": "35ee28f7", + "metadata": {}, + "source": [ + "from tensor_ee.core.errors import UnsupportedOperation\n", + "from tensor_ee.core.grid import Grid\n", + "from tensor_ee.converter.convert import convert\n", + "\n", + "segformer = smp.Segformer(\"mit_b0\", encoder_weights=None, in_channels=2, classes=2).eval()\n", + "try:\n", + " adapt_smp_model(segformer, input_shape=(1, 2, 224, 224))\n", + "except ValueError as refusal:\n", + " print(\"adapt_smp_model:\", refusal)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "01ed7f1c", + "metadata": {}, + "source": [ + "The converter needs the chip's grid. Converting does not compute anything on imagery: it runs\n", + "the model's Python code with placeholder tensors to record its structure. (That also means it\n", + "runs the model's code, so convert only models whose source you trust.)\n", + "\n", + "Its report gives one verdict per building block: **inside** the settings that were validated in\n", + "Earth Engine, or **outside** them." + ] + }, + { + "cell_type": "code", + "id": "43d0b8d0", + "metadata": {}, + "source": [ + "chip = Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000), 224, 224)\n", + "b0 = convert(segformer, (1, 2, 224, 224), chip, name=\"segformer_b0\")\n", + "print(b0.report.lines()[0])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "4fbb3f2c", + "metadata": {}, + "source": [ + "All 18 blocks of SegFormer-B0 are inside their validated records, and the two \"layout\n", + "operations\" only move pixels onto the model's internal grid and back.\n", + "\n", + "SegFormer-B1 has the same structure at larger widths. It converts, but its blocks are\n", + "**outside** their records: the library can express it, and nothing validates it yet. An\n", + "outside verdict is reported, not refused." + ] + }, + { + "cell_type": "code", + "id": "8ead7dab", + "metadata": {}, + "source": [ + "segformer_b1 = smp.Segformer(\"mit_b1\", encoder_weights=None, in_channels=2, classes=2).eval()\n", + "b1 = convert(segformer_b1, (1, 2, 224, 224), chip, name=\"segformer_b1\")\n", + "print(b1.report.lines()[0])\n", + "print(next(node for node in b1.report.nodes if node.id == \"encoder.block1.0\").verdict())" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "82600b1d", + "metadata": {}, + "source": [ + "And a refusal: SegFormer-B0 on a 512 × 512 chip. At that size SegFormer's attention reduces\n", + "its keys and values to a grid larger than one 7 × 7 window, which needs a building block the\n", + "library does not have. The converter says so with a code and the module it stopped at." + ] + }, + { + "cell_type": "code", + "id": "77add088", + "metadata": {}, + "source": [ + "try:\n", + " convert(segformer, (1, 2, 512, 512), Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000), 512, 512))\n", + "except UnsupportedOperation as refusal:\n", + " print(\"Refused:\", refusal.code)\n", + " print(str(refusal)[:400])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "1b23e018", + "metadata": {}, + "source": [ + "## 6. Your model\n", + "\n", + "Put your model in `MY_MODEL` and its chip shape in `MY_INPUT_SHAPE`, then run the next two\n", + "cells. Build the architecture exactly as you trained it, and load your weights if you have\n", + "them (random weights give the same answer). For example:\n", + "\n", + "```python\n", + "MY_MODEL = smp.Unet(\"resnet34\", encoder_weights=None, in_channels=4, classes=3)\n", + "MY_MODEL.load_state_dict(torch.load(\"my_checkpoint.pth\", map_location=\"cpu\", weights_only=True))\n", + "MY_INPUT_SHAPE = (1, 4, 256, 256)\n", + "```\n", + "\n", + "The check runs locally. It tries the CNN path first (for a segmentation-models-pytorch model,\n", + "or a plain PyTorch CNN), and if that does not support the model, the converter for attention\n", + "models." + ] + }, + { + "cell_type": "code", + "id": "5f2159b6", + "metadata": {}, + "source": [ + "MY_MODEL = None # your torch.nn.Module\n", + "MY_INPUT_SHAPE = (1, 2, 224, 224) # (1, bands, height, width)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "382aae7e", + "metadata": {}, + "source": [ + "from tensor_ee.core.errors import UnsupportedOperation\n", + "\n", + "\n", + "def check_my_model(model, input_shape):\n", + " \"\"\"Try the CNN path, then, if that does not support the model, the converter.\"\"\"\n", + " model = model.eval()\n", + " print(\"CNN path (adapt_smp_model and check_compatibility):\")\n", + " try:\n", + " if type(model).__module__.startswith(\"segmentation_models_pytorch\"):\n", + " model_for_check = adapt_smp_model(model, input_shape=input_shape)\n", + " else:\n", + " model_for_check = model # a plain PyTorch CNN is checked as it is\n", + " report = check_compatibility(model_for_check, input_shape)\n", + " explain(report)\n", + " if report.status == \"supported\":\n", + " return report\n", + " except (ValueError, RuntimeError, UnsupportedOperation) as refusal:\n", + " print(\"NOT SUPPORTED:\", refusal)\n", + "\n", + " print()\n", + " print(\"Attention-model path (the converter):\")\n", + " height, width = input_shape[2], input_shape[3]\n", + " try:\n", + " try:\n", + " result = convert(model, input_shape, Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000),\n", + " width, height), name=\"my_model\")\n", + " except UnsupportedOperation as refusal:\n", + " if \"source_grid=None\" not in str(refusal):\n", + " raise\n", + " result = convert(model, input_shape, name=\"my_model\") # a Swin encoder starts on its own grid\n", + " print(result.report.lines()[0])\n", + " print(\"Every block's verdict: print(MY_RESULT.report.text())\")\n", + " return result\n", + " except UnsupportedOperation as refusal:\n", + " print(\"NOT SUPPORTED:\", refusal.code, \"-\", str(refusal)[:600])\n", + " except Exception as failure: # for example, torch.export could not trace the model\n", + " print(\"NOT SUPPORTED: the converter could not read this model:\", type(failure).__name__, failure)\n", + "\n", + "\n", + "if MY_MODEL is None:\n", + " print(\"Put your model in MY_MODEL in the cell above, then run this cell again.\")\n", + "else:\n", + " MY_RESULT = check_my_model(MY_MODEL, MY_INPUT_SHAPE)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "9113a12b", + "metadata": {}, + "source": [ + "## What next\n", + "\n", + "- **Supported, a CNN.** [Notebook 03](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", + " takes a trained segmentation-models-pytorch U-Net from its checkpoint to Earth Engine: it\n", + " plans the tiles, uploads the weights once, predicts one tile, compares it with PyTorch and\n", + " exports a region.\n", + "- **Converted, an attention model.** The\n", + " [converter guide](https://github.com/noelEOS/tensor-ee/blob/main/docs/converter.md) plans the\n", + " model into stages that fit Earth Engine's limits and runs them.\n", + "- **Outside the validated records.** The model can be expressed, but its numbers have no\n", + " evidence yet. Running it and comparing it with PyTorch on your own input is how that evidence\n", + " is made; the maintainer can advise.\n", + "- **Refused.** See the options under step 4.\n", + "\n", + "The [block catalogue](https://github.com/noelEOS/tensor-ee/blob/main/docs/building-blocks.md)\n", + "lists every building block with the settings it was validated at." + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/03_bring_your_own_cnn.ipynb b/notebooks/03_bring_your_own_cnn.ipynb new file mode 100644 index 0000000..01fbc39 --- /dev/null +++ b/notebooks/03_bring_your_own_cnn.ipynb @@ -0,0 +1,693 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "ab4596ab", + "metadata": {}, + "source": [ + "# Bring your own CNN: a trained U-Net from checkpoint to Earth Engine\n", + "\n", + "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", + "\n", + "This notebook takes a U-Net you trained with\n", + "[segmentation-models-pytorch](https://github.com/qubvel-org/segmentation_models.pytorch)\n", + "(SMP) and runs it inside Earth Engine on your own imagery. The steps:\n", + "\n", + "1. load your checkpoint and check that TensorEE supports the model;\n", + "2. work out how much extra border (the halo) each tile needs, and plan the tiles;\n", + "3. upload the model's weights to Earth Engine once, as table assets;\n", + "4. predict one tile and compare it with PyTorch on the same pixels;\n", + "5. export a whole region.\n", + "\n", + "**Offline first.** With the settings as they are, the notebook uses a small demonstration\n", + "U-Net with random weights and runs every step offline: it builds the Earth Engine\n", + "computations but sends nothing. Then you set your checkpoint, your project and your input,\n", + "and turn on the billed steps one at a time. Each billed step says beforehand what it creates\n", + "and what is known of its cost, and runs only when both `RUN_LIVE` and its own switch are\n", + "`True`.\n", + "\n", + "**What has been checked.** Each building block this U-Net uses agrees with PyTorch at the\n", + "settings it was validated at. One trained U-Net with a ResNet-18 encoder, prepared the way\n", + "this notebook prepares yours, ran in Earth Engine on a 64 × 64 test input and matched PyTorch\n", + "within 1.2e-5 in the logits. The tiled recipe in this notebook has **not** been run in Earth\n", + "Engine as a whole: step 9's comparison is how you check it for your model. Agreement with\n", + "PyTorch says nothing about accuracy, which you measure on your own labelled data." + ] + }, + { + "cell_type": "markdown", + "id": "b29a1f17", + "metadata": {}, + "source": [ + "## 1. Install\n", + "\n", + "The `segwater` extra brings segmentation-models-pytorch 0.5.0, the version the library was\n", + "checked with. PyTorch is already installed in Colab." + ] + }, + { + "cell_type": "code", + "id": "b0fec902", + "metadata": {}, + "source": [ + "%pip install --quiet \"tensor-ee[segwater]\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "2f178b02", + "metadata": {}, + "source": [ + "## 2. Settings\n", + "\n", + "Edit this cell, and only this cell.\n", + "\n", + "**Your model.** `CHECKPOINT` is the path of your trained weights, saved with\n", + "`torch.save(model.state_dict(), path)`. In Colab, upload it with the **Files** panel on the left.\n", + "Leave it `None` to use the small demonstration model, which runs offline only. `ENCODER`,\n", + "`BANDS` and `CLASSES` must match how you trained the model.\n", + "\n", + "**Your input.** `INPUT_IMAGE` is an Earth Engine image with the bands `BANDS`, already scaled\n", + "the way your training data was (TensorEE does not normalize for you). `CRS`, `REGION_X`,\n", + "`REGION_Y` and `PIXEL` place the region: the upper-left corner of its first tile, in metres of\n", + "`CRS`. The image must also cover a border of one halo around the region (step 5 prints how\n", + "wide).\n", + "\n", + "**The billed steps.** Each runs only when `RUN_LIVE` and its own switch are both `True`.\n", + "\n", + "After changing `RUN_LIVE`, restart the runtime (in Colab: **Runtime → Restart session**) and\n", + "run the notebook again from the top: the offline mode stays installed until the runtime\n", + "restarts." + ] + }, + { + "cell_type": "code", + "id": "6193d146", + "metadata": {}, + "source": [ + "PROJECT = \"YOUR_EE_PROJECT\" # your Cloud project registered for Earth Engine\n", + "RUN_LIVE = False # True: sign in; the switches below then choose the billed steps\n", + "\n", + "# Your model\n", + "CHECKPOINT = None # for example \"my_unet.pth\"; None = demonstration model, offline only\n", + "ENCODER = \"resnet34\"\n", + "BANDS = (\"VV\", \"VH\") # your input bands, in training order\n", + "CLASSES = 2\n", + "\n", + "# Your input and the region\n", + "INPUT_IMAGE = \"projects/YOUR_EE_PROJECT/assets/YOUR_PREPARED_INPUT\"\n", + "CRS = \"EPSG:32651\"\n", + "REGION_X, REGION_Y = 286030, 2640860 # upper-left corner of the region, in CRS metres\n", + "PIXEL = 10 # metres per pixel\n", + "TILE = 1024 # tile side for the region, in pixels\n", + "REGION_COLUMNS, REGION_ROWS = 2, 1 # the region is this many tiles\n", + "TEST_TILE = 256 # side of the single test tile of step 8\n", + "\n", + "# What the live steps create\n", + "WEIGHTS_MANIFEST = \"my_unet_weights.json\" # a local file recording the uploaded weights\n", + "TEST_ASSET = \"my_unet_test_tile\" # new image assets in your project\n", + "REGION_ASSET = \"my_unet_region\"\n", + "\n", + "# The billed steps\n", + "UPLOAD_WEIGHTS = False # step 7\n", + "PREDICT_TEST_TILE = False # step 8\n", + "COMPARE_TEST_TILE = False # step 9\n", + "EXPORT_REGION = False # step 10" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "2cfc7c5f", + "metadata": {}, + "source": [ + "PLACEHOLDER = \"YOUR_EE_PROJECT\"\n", + "\n", + "\n", + "def live_step(what, switch=True):\n", + " \"\"\"True when a billed step should run; refuses a placeholder project or a missing checkpoint.\"\"\"\n", + " if not (RUN_LIVE and switch):\n", + " reason = \"RUN_LIVE is False\" if not RUN_LIVE else \"its switch is False\"\n", + " print(f\"Skipped: {what}. {reason}, so nothing was sent to Earth Engine.\")\n", + " return False\n", + " if not PROJECT or PROJECT == PLACEHOLDER:\n", + " raise ValueError(\"Set PROJECT to your own Cloud project ID in the settings cell.\")\n", + " if CHECKPOINT is None:\n", + " raise ValueError(\"Live steps need your trained checkpoint: set CHECKPOINT. \"\n", + " \"The demonstration model has random weights.\")\n", + " return True" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "695f4df2", + "metadata": {}, + "source": [ + "## 3. Sign in, or work offline\n", + "\n", + "With `RUN_LIVE = True` the first cell signs you in; the first time, `ee.Authenticate()` shows a\n", + "link or a prompt. Signing in costs nothing by itself. With `RUN_LIVE = False` the second cell\n", + "installs the Earth Engine client's offline function signatures (`ee.apitestcase.ApiTestCase`,\n", + "a test helper that ships with the client), so the notebook can build Earth Engine computations\n", + "without credentials and without sending anything." + ] + }, + { + "cell_type": "code", + "id": "84bcd7d7", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "import ee\n", + "\n", + "if live_step(\"connect to Earth Engine\"):\n", + " ee.Authenticate()\n", + " ee.Initialize(project=PROJECT)\n", + " print(\"Signed in. Requests are billed to\", PROJECT)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "d4b2bff3", + "metadata": {}, + "source": [ + "import ee\n", + "\n", + "if not RUN_LIVE:\n", + " from ee.apitestcase import ApiTestCase\n", + "\n", + " offline = ApiTestCase()\n", + " offline.setUp() # offline signatures: no credentials, no network\n", + " print(\"Offline: the notebook builds Earth Engine computations but sends none.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "aa4d0e83", + "metadata": {}, + "source": [ + "## 4. Load the model\n", + "\n", + "With `CHECKPOINT = None` this builds a shallow U-Net (three encoder stages, small decoder)\n", + "with random weights, so that the offline steps stay quick. With your checkpoint, it builds the\n", + "standard U-Net and loads your weights strictly: a missing or unexpected weight name stops it,\n", + "which usually means the settings do not match your training." + ] + }, + { + "cell_type": "code", + "id": "42d821f2", + "metadata": {}, + "source": [ + "import torch\n", + "import segmentation_models_pytorch as smp\n", + "\n", + "torch.set_num_threads(2)\n", + "\n", + "\n", + "def make_model(**changes):\n", + " settings = dict(encoder_name=ENCODER, encoder_weights=None, in_channels=len(BANDS),\n", + " classes=CLASSES, activation=None)\n", + " settings.update(changes)\n", + " return smp.Unet(**settings).eval()\n", + "\n", + "\n", + "if CHECKPOINT is None:\n", + " torch.manual_seed(20261020)\n", + " model = make_model(encoder_depth=3, decoder_channels=(16, 8, 4))\n", + " tile, test_tile = 64, 64 # small tiles for the demonstration\n", + " print(\"Demonstration model: a shallow U-Net with random weights.\")\n", + "else:\n", + " model = make_model()\n", + " model.load_state_dict(torch.load(CHECKPOINT, map_location=\"cpu\", weights_only=True), strict=True)\n", + " tile, test_tile = TILE, TEST_TILE\n", + " print(\"Loaded\", CHECKPOINT)\n", + "print(f\"{sum(p.numel() for p in model.parameters()):,} parameters\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "73b9c3ce", + "metadata": {}, + "source": [ + "## 5. Check support, and find the halo\n", + "\n", + "`adapt_smp_model` prepares the model for TensorEE at a fixed input size, and checks that the\n", + "prepared model gives exactly the same numbers as yours on a test input. `check_compatibility`\n", + "then checks every operation. (Notebook 02 explains these reports in more detail.)\n", + "\n", + "`analyze_receptive_field` works out how far each output pixel can see into the input. That\n", + "reach is the **halo**: to compute a tile exactly as the model would compute it on a larger\n", + "image, Earth Engine must read the tile plus a halo on every side, and the halo is thrown away\n", + "afterwards. For a standard U-Net with a ResNet-34 encoder the halo is 544 pixels, so a\n", + "1,024-pixel tile reads 2,112 × 2,112 pixels: about 4.25 input pixels per output pixel. Larger\n", + "tiles waste less." + ] + }, + { + "cell_type": "code", + "id": "3c1fd62e", + "metadata": {}, + "source": [ + "from tensor_ee import check_compatibility\n", + "from tensor_ee.receptive_field import analyze_receptive_field\n", + "from tensor_ee.smp import adapt_smp_model\n", + "\n", + "\n", + "def prepare(height, width):\n", + " \"\"\"The model adapted to one input size, refused unless every operation is supported.\"\"\"\n", + " shape = (1, len(BANDS), height, width)\n", + " adapted = adapt_smp_model(model, input_shape=shape)\n", + " report = check_compatibility(adapted, shape)\n", + " if report.status != \"supported\":\n", + " raise ValueError(f\"TensorEE does not support this model: {report.to_dict()['blockers'][:3]}\")\n", + " return adapted\n", + "\n", + "\n", + "probe = prepare(64, 64)\n", + "support = analyze_receptive_field(probe)\n", + "halo, alignment = support.safe_halo, support.alignment\n", + "print(\"Supported. Prepared model equals the original on a test input:\",\n", + " probe.tensor_ee_smp_metadata[\"torch_equal\"])\n", + "print(f\"Halo: {halo} pixels on each side ({halo * PIXEL / 1000:g} km). \"\n", + " f\"Tile sizes and positions must be multiples of {alignment} pixels.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "41ea3bf8", + "metadata": {}, + "source": [ + "## 6. Plan the tiles\n", + "\n", + "The plan puts the region's tiles side by side and gives each one its halo. Every tile then\n", + "reads an input of the same size, which the model needs. The region's input must cover the\n", + "region plus the halo on every side.\n", + "\n", + "There are two plans: one test tile, for step 8, and the region, for step 10." + ] + }, + { + "cell_type": "code", + "id": "1a50cc0b", + "metadata": {}, + "source": [ + "from tensor_ee import Grid\n", + "from tensor_ee.converter import PixelWindow, plan_tiles\n", + "\n", + "\n", + "def plan_region(side, columns, rows):\n", + " \"\"\"Tiles of side x side pixels, with a full halo of real input around the whole region.\"\"\"\n", + " for name, value in ((\"tile side\", side), (\"halo\", halo)):\n", + " if value % alignment:\n", + " raise ValueError(f\"The {name} ({value}) must be a multiple of {alignment} pixels.\")\n", + " width, height = columns * side + 2 * halo, rows * side + 2 * halo\n", + " padded = Grid(CRS, (PIXEL, 0, REGION_X - halo * PIXEL, 0, -PIXEL, REGION_Y + halo * PIXEL),\n", + " width, height)\n", + " region = PixelWindow(halo, halo, columns * side, rows * side)\n", + " return plan_tiles(padded, side, halo, alignment=alignment, roi=region)\n", + "\n", + "\n", + "def describe(plan, name):\n", + " context = plan.tiles[0].context\n", + " read = sum(t.context.pixels for t in plan.tiles)\n", + " print(f\"{name}: {len(plan.tiles)} tile(s) of {plan.core_size[0]} x {plan.core_size[1]} pixels; \"\n", + " f\"each reads {context.width} x {context.height}; \"\n", + " f\"{read / plan.roi.pixels:.2f} input pixels per output pixel.\")\n", + " print(f\" Input needed: {plan.global_grid.width} x {plan.global_grid.height} pixels \"\n", + " f\"from corner {plan.global_grid.transform[2]}, {plan.global_grid.transform[5]} in {CRS}.\")\n", + "\n", + "\n", + "test_plan = plan_region(test_tile, 1, 1)\n", + "region_plan = plan_region(tile, REGION_COLUMNS, REGION_ROWS)\n", + "describe(test_plan, \"Test tile\")\n", + "describe(region_plan, \"Region\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "47fb0596", + "metadata": {}, + "source": [ + "## 7. Upload the weights once (billed to your project)\n", + "\n", + "TensorEE folds each BatchNorm layer into the convolution before it (the same arithmetic, fewer\n", + "steps) and stores the resulting numbers in Earth Engine table assets. Every prediction then\n", + "reads them from there, so they are uploaded once and reused.\n", + "\n", + "**What this creates.** Table assets in your project, under a new name\n", + "`projects/PROJECT/assets/tensor_ee_weights_`, each holding at most 8 MB of numbers as\n", + "text; the cell prints an estimate of how many. One batch table-export task per table, at most\n", + "two at a time; the cell waits for all of them. Locally, the file `WEIGHTS_MANIFEST`, which\n", + "records the tables and a fingerprint of every convolution. Keep that file: the later steps\n", + "need it, and in Colab it disappears when the runtime ends, so download it from the **Files**\n", + "panel.\n", + "\n", + "**What it costs.** Not recorded. The trained ResNet-18 U-Net of notebook 01 has 14.3 million\n", + "numbers and took 29 tables." + ] + }, + { + "cell_type": "code", + "id": "d6897497", + "metadata": {}, + "source": [ + "import math\n", + "\n", + "from tensor_ee import effective_convolutions\n", + "\n", + "parameters = effective_convolutions(probe)\n", + "values = sum(p.weights.size + (0 if p.bias is None else p.bias.size) for p in parameters)\n", + "print(f\"{len(parameters)} folded convolutions, {values:,} numbers: \"\n", + " f\"about {math.ceil(values * 14 / 7_000_000)} tables.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "1aaa8523", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "import json\n", + "import uuid\n", + "from collections import Counter\n", + "from pathlib import Path\n", + "\n", + "from tensor_ee import export_weights\n", + "\n", + "if live_step(\"upload the weights\", UPLOAD_WEIGHTS):\n", + " manifest_path = Path(WEIGHTS_MANIFEST)\n", + " if manifest_path.exists():\n", + " raise FileExistsError(f\"{manifest_path} already exists: the weights are uploaded. \"\n", + " \"Use a new WEIGHTS_MANIFEST name only for new weights.\")\n", + " namespace = f\"projects/{PROJECT}/assets/tensor_ee_weights_{uuid.uuid4().hex[:12]}\"\n", + "\n", + " def progress(manifest):\n", + " states = Counter(shard[\"status\"] for shard in manifest[\"shards\"])\n", + " print(json.dumps({\"status\": manifest[\"status\"], \"tables\": dict(states)}), flush=True)\n", + "\n", + " export_weights(parameters, namespace, manifest_path, on_update=progress)\n", + " print(\"Weights uploaded; manifest written to\", manifest_path)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "09c5f873", + "metadata": {}, + "source": [ + "## 8. Predict the test tile (billed to your project)\n", + "\n", + "The next cell builds the prediction: the model, compiled to Earth Engine operations for the test tile's input size, with\n", + "its weights read from the uploaded tables. Offline, the tables are placeholders that do not\n", + "exist, which is enough to build and measure the computation.\n", + "\n", + "The result has these bands: `score_0`, `score_1`, … (the raw logits), `p_0`, `p_1`, …\n", + "(probabilities), `class_index` (the most likely class) and, for two classes, `binary_class`\n", + "(1 where `p_1` ≥ 0.5). A prediction is made only where every input band is present." + ] + }, + { + "cell_type": "code", + "id": "0c458253", + "metadata": {}, + "source": [ + "import json\n", + "from pathlib import Path\n", + "\n", + "from tensor_ee import MatchedWeightProvider, Raster, binary_threshold, segmentation_outputs\n", + "from tensor_ee.converter import compile_tiled_model\n", + "from tensor_ee.converter.request import deep_recursion\n", + "\n", + "SCORES = tuple(f\"score_{k}\" for k in range(CLASSES))\n", + "PROBABILITIES = tuple(f\"p_{k}\" for k in range(CLASSES))\n", + "\n", + "\n", + "def weights_manifest():\n", + " \"\"\"The uploaded weights, or offline a stand-in that names tables which do not exist.\"\"\"\n", + " if RUN_LIVE and Path(WEIGHTS_MANIFEST).exists():\n", + " return json.loads(Path(WEIGHTS_MANIFEST).read_text())\n", + " if RUN_LIVE:\n", + " raise FileNotFoundError(f\"{WEIGHTS_MANIFEST} not found: upload the weights first (step 7).\")\n", + " return {\"format\": \"native_ee_featurecollection_weights_v2\", \"status\": \"ready\",\n", + " \"layers\": {p.key: {\"shape\": list(p.shape), \"dtype\": p.dtype, \"bias\": p.bias is not None,\n", + " \"source\": p.source, \"sha256\": p.sha256,\n", + " \"assets\": [\"projects/YOUR_EE_PROJECT/assets/weights_not_uploaded\"]}\n", + " for p in parameters}}\n", + "\n", + "\n", + "def input_image():\n", + " if RUN_LIVE:\n", + " return ee.Image(INPUT_IMAGE)\n", + " return ee.Image.constant([0] * len(BANDS)).rename(list(BANDS)) # offline stand-in\n", + "\n", + "\n", + "def build_prediction(plan):\n", + " \"\"\"The model's Earth Engine computation over a plan's tiles, with its output bands.\"\"\"\n", + " context = plan.tiles[0].context\n", + " adapted = prepare(context.height, context.width)\n", + " if analyze_receptive_field(adapted).safe_halo != halo:\n", + " raise ValueError(\"The halo changed with the input size; revise the plan.\")\n", + " provider = MatchedWeightProvider(adapted, weights_manifest())\n", + " source = input_image().select(list(BANDS)).toFloat().reproject(plan.global_grid.projection())\n", + " valid = Raster(source.mask().reduce(ee.Reducer.min()).rename(\"valid\"), (\"valid\",), plan.global_grid)\n", + " logits = compile_tiled_model(adapted, source, BANDS, plan, weight_provider=provider)\n", + " outputs = segmentation_outputs(logits, logit_bands=SCORES, probability_bands=PROBABILITIES,\n", + " validity=valid)\n", + " bands = [outputs.logits.image, outputs.probabilities.image, outputs.classes.image]\n", + " if CLASSES == 2:\n", + " bands.append(binary_threshold(outputs.probabilities, positive_class=1, threshold=0.5,\n", + " output_band=\"binary_class\").image)\n", + " return ee.Image.cat(bands).toFloat() # one type for export; class values stay exact\n", + "\n", + "\n", + "def export_task(image, plan, asset_name, description):\n", + " grid = plan.roi_grid\n", + " return ee.batch.Export.image.toAsset(\n", + " image=image, description=description, assetId=f\"projects/{PROJECT}/assets/{asset_name}\",\n", + " region=grid.geometry(), crs=grid.crs, crsTransform=list(grid.transform),\n", + " maxPixels=grid.width * grid.height, pyramidingPolicy={\".default\": \"sample\"}, overwrite=False)\n", + "\n", + "\n", + "with deep_recursion():\n", + " test_prediction = build_prediction(test_plan)\n", + " test_graph_bytes = len(test_prediction.serialize().encode())\n", + "print(f\"Test tile computation: {test_graph_bytes:,} bytes of Earth Engine expression.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "757fef81", + "metadata": {}, + "source": [ + "**What the export creates.** One batch export task and, when it completes, one new image\n", + "asset `projects/PROJECT/assets/TEST_ASSET` over the test tile (not its halo).\n", + "\n", + "**What it costs.** Not measured for this recipe. For scale: the trained water U-Net's\n", + "224 × 224 chip took 64.9 batch EECU-seconds, and this test tile reads a larger input\n", + "(step 6). If Earth Engine reports that it ran out of memory or time, reduce `TEST_TILE` and\n", + "try again; do not just resubmit." + ] + }, + { + "cell_type": "code", + "id": "df5fae6a", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "import time\n", + "\n", + "\n", + "def wait(task, minutes=60):\n", + " \"\"\"Check a task every 30 seconds until it ends.\"\"\"\n", + " deadline = time.time() + minutes * 60\n", + " while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", + " if time.time() > deadline:\n", + " raise TimeoutError(\"Still running; run this cell again later, or watch the Tasks tab.\")\n", + " time.sleep(30)\n", + " print(\"Task\", task.id, state)\n", + " if state != \"COMPLETED\":\n", + " raise RuntimeError(task.status().get(\"error_message\", \"the export did not complete\"))\n", + "\n", + "\n", + "if live_step(\"export the test tile (1 batch task, 1 new asset)\", PREDICT_TEST_TILE):\n", + " test_task = export_task(test_prediction, test_plan, TEST_ASSET, \"tensor_ee_test_tile\")\n", + " test_task.start()\n", + " print(\"Started\", test_task.id)\n", + " wait(test_task)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "3d569881", + "metadata": {}, + "source": [ + "## 9. Compare the test tile with PyTorch (billed to your project)\n", + "\n", + "This is the check that matters for your model: Earth Engine's logits against your model's in\n", + "PyTorch, on exactly the same input pixels. The cell downloads the test tile's input (tile plus\n", + "halo) and the exported logits, runs your model on the input locally, and reports the largest\n", + "absolute difference. TensorEE's target is 1e-4.\n", + "\n", + "**What this does.** Two interactive requests (`ee.data.computePixels`): the input and the\n", + "exported result. It creates nothing. The input must be complete over the tile and its halo.\n", + "\n", + "**What it costs.** Small interactive requests, not measured. The cell refuses an input over\n", + "40 MB." + ] + }, + { + "cell_type": "code", + "id": "b3f6e6aa", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "import numpy as np\n", + "\n", + "\n", + "def download(image, grid, bands):\n", + " \"\"\"The exact float32 pixels of an image's bands on a grid: one interactive request.\"\"\"\n", + " names = (\"scaleX\", \"shearX\", \"translateX\", \"shearY\", \"scaleY\", \"translateY\")\n", + " values = ee.data.computePixels({\n", + " \"expression\": image.select(list(bands)), \"fileFormat\": \"NUMPY_NDARRAY\",\n", + " \"grid\": {\"dimensions\": {\"width\": grid.width, \"height\": grid.height},\n", + " \"affineTransform\": dict(zip(names, grid.transform)), \"crsCode\": grid.crs}})\n", + " return np.stack([values[band] for band in bands]).astype(np.float32)\n", + "\n", + "\n", + "if live_step(\"compare the test tile with PyTorch (2 requests)\", COMPARE_TEST_TILE):\n", + " padded = test_plan.global_grid\n", + " size = padded.width * padded.height * (len(BANDS) + 1) * 4\n", + " if size > 40_000_000:\n", + " raise ValueError(f\"The input is {size / 1e6:.0f} MB; reduce TEST_TILE.\")\n", + " source = ee.Image(INPUT_IMAGE).select(list(BANDS)).toFloat().reproject(padded.projection())\n", + " valid = source.mask().reduce(ee.Reducer.min()).rename(\"valid\")\n", + " pixels = download(source.unmask(0).addBands(valid), padded, (*BANDS, \"valid\"))\n", + " if pixels[-1].min() < 1:\n", + " raise ValueError(\"The input has missing pixels in the tile or its halo; the comparison \"\n", + " \"needs complete input. Move the region or fill the gaps.\")\n", + " exported = download(ee.Image(f\"projects/{PROJECT}/assets/{TEST_ASSET}\"), test_plan.roi_grid, SCORES)\n", + "\n", + " with torch.no_grad():\n", + " reference = model(torch.from_numpy(pixels[None, :-1].copy()))[0].numpy()\n", + " core = reference[:, halo:halo + test_tile, halo:halo + test_tile]\n", + " error = np.abs(exported - core)\n", + " disagreements = int((exported.argmax(0) != core.argmax(0)).sum())\n", + " print(f\"Largest absolute logit difference: {error.max():.3e} (target 1e-4)\")\n", + " print(f\"Mean absolute difference: {error.mean():.3e}; \"\n", + " f\"class disagreements: {disagreements} of {core[0].size} pixels\")\n", + " print(\"Within the target.\" if error.max() <= 1e-4 else\n", + " \"Above the target: please report it, with these numbers, as an issue.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "0713bc1b", + "metadata": {}, + "source": [ + "## 10. Export the region (billed to your project)\n", + "\n", + "**What this creates.** One batch export task and, when it completes, one new image asset\n", + "`projects/PROJECT/assets/REGION_ASSET` over the region (not its halo), with the bands of\n", + "step 8, in float32.\n", + "\n", + "**What it costs.** Not measured. Earth Engine computes every tile with its halo (step 6\n", + "prints how many input pixels it reads per output pixel), and capacity at these tile sizes has\n", + "not been measured: an export can fail for lack of memory. Run the test tile first." + ] + }, + { + "cell_type": "code", + "id": "15be8dc7", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "if live_step(\"export the region (1 batch task, 1 new asset)\", EXPORT_REGION):\n", + " with deep_recursion():\n", + " region_prediction = build_prediction(region_plan)\n", + " region_task = export_task(region_prediction, region_plan, REGION_ASSET, \"tensor_ee_region\")\n", + " region_task.start()\n", + " print(\"Started\", region_task.id, \"- follow it in the Tasks tab, or wait here.\")\n", + " wait(region_task, minutes=240)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "4cd0dc7c", + "metadata": {}, + "source": [ + "## What this shows, and what it does not\n", + "\n", + "- **Offline**, the notebook shows that TensorEE supports the model and can build its Earth\n", + " Engine computation for your tiles.\n", + "- **Step 9** is live agreement for your model, on one tile and one input. It is the evidence\n", + " to keep, with its numbers.\n", + "- **The region** reuses the same computation for every tile; the halo is what makes tile\n", + " borders match a single run of the model, as far as the model's reach goes.\n", + "- **Accuracy** is a separate question: evaluate the exported predictions against your own\n", + " labelled data.\n", + "\n", + "The [guide this notebook follows](https://github.com/noelEOS/tensor-ee/blob/main/docs/own-model.md)\n", + "gives the details: other encoders, irregular regions, and how validity masks are handled.\n", + "**Cleaning up**: the weight tables and image assets stay in your project until you delete them\n", + "in the Code Editor's **Assets** tab; keep the weight tables as long as you predict with them." + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py new file mode 100644 index 0000000..08ddb33 --- /dev/null +++ b/tests/test_notebooks.py @@ -0,0 +1,187 @@ +"""The reader notebooks in ``notebooks/``: clean, free of private values, and runnable offline. + +Read with ``json`` (no Jupyter needed). Each notebook's code cells run in order in a fresh +namespace with ``RUN_LIVE = False``, skipping the cells tagged ``live`` and the ``%`` and ``!`` +lines (installs). Every Earth Engine cloud entry point is refused while they run, as in +``demos.code_editor.build.offline_earth_engine``. Notebooks that need segmentation-models-pytorch +or timm skip by name without them. +""" +from __future__ import annotations + +import contextlib +import json +import re +import sys +from pathlib import Path +from unittest import mock + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +NOTEBOOKS = ROOT / "notebooks" +EXPECTED = { + "01_quickstart.ipynb": (), + "02_is_my_model_supported.ipynb": ("segmentation_models_pytorch", "timm"), + "03_bring_your_own_cnn.ipynb": ("segmentation_models_pytorch",), +} +PLACEHOLDER_PROJECT = "YOUR_EE_PROJECT" +MANIFESTS = ROOT / "data" / "manifests" +COLAB = "https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/" +# Calls that reach Earth Engine or start work: only cells tagged "live" may make them. +LIVE_CALLS = re.compile(r"ee\.Initialize|ee\.Authenticate|\.getInfo\(|\.start\(\)|computePixels|" + r"getMapId|\.status\(\)|export_weights\(") +EMAIL = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+\.[A-Za-z0-9.-]+") +HOME = re.compile(r"/Users/|/home/|[A-Za-z]:\\\\Users|~/") +ASSET = re.compile(r"projects/([^/\s\"'`]+)/assets/[^\s\"'`)]*") +PROJECT_VALUE = re.compile(r"""project\s*=\s*["']([^"']*)["']""", re.I) + + +def _cells(name): + return json.loads((NOTEBOOKS / name).read_text())["cells"] + + +def _source(cell): + return "".join(cell["source"]) + + +def _public_assets(): + assets = set() + for path in MANIFESTS.glob("*.json"): + def walk(value): + if isinstance(value, dict): + for item in value.values(): + walk(item) + elif isinstance(value, list): + for item in value: + walk(item) + elif isinstance(value, str) and "/assets/" in value: + assets.add(value) + walk(json.loads(path.read_text())) + return assets + + +def test_the_notebooks_are_the_three_reader_notebooks(): + assert {path.name for path in NOTEBOOKS.glob("*.ipynb")} == set(EXPECTED) + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_outputs_and_execution_counts_are_cleared(name): + notebook = json.loads((NOTEBOOKS / name).read_text()) + assert notebook["nbformat"] == 4 + for cell in notebook["cells"]: + assert cell["cell_type"] in ("markdown", "code") + if cell["cell_type"] == "code": + assert cell["outputs"] == [] and cell["execution_count"] is None, _source(cell)[:80] + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_each_opens_with_a_colab_badge_for_its_file_on_main(name): + first = _source(_cells(name)[0]) + assert first.startswith("# ") and f"({COLAB}{name})" in first + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_the_settings_start_offline_and_only_live_cells_reach_earth_engine(name): + cells = [cell for cell in _cells(name) if cell["cell_type"] == "code"] + settings = [cell for cell in cells if re.search(r"^RUN_LIVE = False\b", _source(cell), re.M)] + assert len(settings) == 1, "one settings cell with RUN_LIVE = False" + live = [cell for cell in cells if "live" in cell["metadata"].get("tags", [])] + if live: + assert f'PROJECT = "{PLACEHOLDER_PROJECT}"' in _source(settings[0]) + for cell in cells: + source = _source(cell) + if cell in live: + assert "live_step(" in source, "a live cell checks RUN_LIVE and the project first" + else: + assert not LIVE_CALLS.search(source), source[:200] + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_no_private_values(name): + text = "\n".join(_source(cell) for cell in _cells(name)) + assert not EMAIL.search(text) + assert not HOME.search(text) + public = _public_assets() + for match in ASSET.finditer(text): + project = match.group(1) + # The placeholder, the settings variable, or a public weight table. + assert project in (PLACEHOLDER_PROJECT, "{PROJECT}", "PROJECT") or match.group(0) in public, match.group(0) + for value in PROJECT_VALUE.findall(text): + assert value == PLACEHOLDER_PROJECT, value + + +@contextlib.contextmanager +def _offline(): + """Refuse every Earth Engine data call and authentication; restore the client afterwards.""" + import ee + from demos.code_editor.build import CloudCall, _REFUSED_DATA_CALLS + + def refuse(name): + def call(*args, **kwargs): + raise CloudCall(f"cloud call attempted while running a notebook offline: {name}") + return call + + install = ee.data._install_cloud_api_resource + limit, path = sys.getrecursionlimit(), list(sys.path) + try: + with contextlib.ExitStack() as stack: + for name in _REFUSED_DATA_CALLS: + if hasattr(ee.data, name): + stack.enter_context(mock.patch.object(ee.data, name, refuse(name))) + # ee.Initialize stays: the notebooks' ApiTestCase calls it with its offline client. + # The static check above keeps it out of every cell that is not tagged live. + stack.enter_context(mock.patch.object(ee, "Authenticate", refuse("Authenticate"))) + yield + finally: + ee.data._install_cloud_api_resource = install + ee.Reset() + sys.setrecursionlimit(limit) + sys.path[:] = path + + +def _run(name, monkeypatch): + for module in EXPECTED[name]: + pytest.importorskip(module) + import torch + monkeypatch.chdir(ROOT) + namespace = {"__name__": "__notebook__"} + threads = torch.get_num_threads() + with _offline(): + try: + for index, cell in enumerate(_cells(name)): + if cell["cell_type"] != "code" or "live" in cell["metadata"].get("tags", []): + continue + lines = [line for line in _source(cell).split("\n") + if not line.lstrip().startswith(("%", "!"))] + exec(compile("\n".join(lines), f"{name}[{index}]", "exec"), namespace) + if "RUN_LIVE" in namespace: + assert namespace["RUN_LIVE"] is False + finally: + if "offline" in namespace: + namespace["offline"].tearDown() + torch.set_num_threads(threads) + return namespace + + +def test_quickstart_runs_offline_and_builds_the_recorded_graph(monkeypatch): + import ee + namespace = _run("01_quickstart.ipynb", monkeypatch) + assert namespace["OUTPUT_BANDS"] == ["logit_0", "logit_1", "water_probability", "water_mask", "coverage"] + assert namespace["scene_id"] == namespace["PRESET"]["scene"] + if ee.__version__ == namespace["METADATA"]["earthengineApiVersion"]: + # predict_segwater's graph for the preset chip, built with the trained checkpoint. + assert namespace["graph_sha256"] == namespace["preset"]["library_sha256"] + + +def test_is_my_model_supported_runs_offline(monkeypatch): + namespace = _run("02_is_my_model_supported.ipynb", monkeypatch) + assert namespace["MY_MODEL"] is None + assert namespace["b0"].report.lines()[0].startswith("segformer_b0: 20 nodes; 18 with a validated record, 18 inside") + + +def test_bring_your_own_cnn_runs_offline(monkeypatch): + namespace = _run("03_bring_your_own_cnn.ipynb", monkeypatch) + assert namespace["CHECKPOINT"] is None + assert (namespace["halo"], namespace["alignment"]) == (112, 8) + assert len(namespace["region_plan"].tiles) == 2 and len(namespace["test_plan"].tiles) == 1 + assert namespace["test_graph_bytes"] > 0 From 29e2873215741e6c4270f6e9f3c7e3a4c715a5ec Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:26:10 +0800 Subject: [PATCH 2/9] Make notebook 01 a true quickstart and tighten the reader notebooks Notebook 01 now has one setting, the project, in a Colab form; an empty project stops with a one-line message. Install, settings, sign in, place and scene, build, export and map are one cell each, with the background in three short bullets at the end. It no longer has an offline mode, a restart instruction or the byte-identity check; the test supplies the offline signatures, the project and the preset place itself, and keeps the identity assertion. The asset name carries a timestamp, so a rerun does not collide. Notebook 03 undoes its offline signatures when RUN_LIVE is switched on, so no restart is needed, and the test checks it. Notebooks 02 and 03 and docs/start-here.md use bullets and say that runs use the project's Earth Engine quota: free on noncommercial projects, the plan's rate on commercial ones; the quickstart's export used about 65 EECU-seconds. --- docs/start-here.md | 148 +++--- notebooks/01_quickstart.ipynb | 627 +++++------------------ notebooks/02_is_my_model_supported.ipynb | 159 +++--- notebooks/03_bring_your_own_cnn.ipynb | 463 ++++++----------- tests/test_notebooks.py | 139 +++-- 5 files changed, 524 insertions(+), 1012 deletions(-) diff --git a/docs/start-here.md b/docs/start-here.md index 0e8d692..930dc4c 100644 --- a/docs/start-here.md +++ b/docs/start-here.md @@ -1,101 +1,75 @@ # Start here -This page is for anyone with an Earth Engine account and some Python who wants to try -TensorEE. You do not need to know how the library is built, and you do not need any special -tools: a browser, and for the notebooks, Google Colab or Jupyter. +For anyone with an Earth Engine account and some Python who wants to try TensorEE. A browser +is enough: the notebooks run in Google Colab. -## What TensorEE does, and what it does not +## What TensorEE does, and does not do -TensorEE runs a trained neural network inside Earth Engine. You keep your imagery in Earth -Engine, and the model's arithmetic (convolutions, additions, activations) is written as -ordinary Earth Engine image and array operations, so Earth Engine computes the prediction -itself. There is no model server to deploy and no imagery to download. The weights are stored -as Earth Engine table assets. It works only for the kinds of model it supports: some -convolutional networks (CNNs) such as standard U-Nets, and a few attention models through a -converter. It refuses anything else with a message saying what is missing, rather than -running an approximation. It does not train models, and it does not tell you how accurate a -model is: it checks that Earth Engine computes the same numbers as PyTorch, which is a -different question from whether the predictions are right. +- **Does:** runs a trained neural network inside Earth Engine. The model's arithmetic becomes + ordinary Earth Engine image and array operations, so Earth Engine computes the prediction + where your imagery already is. No model server, no image download; the weights are Earth + Engine table assets. +- **Supports:** some CNNs, such as standard U-Nets, and a few attention models through a + converter. Anything else is refused with a message naming what is missing, never + approximated. +- **Does not:** train models, or tell you how accurate one is. It checks that Earth Engine + computes the same numbers as PyTorch, which is a different question. ## Pick a path -### 1. See it work, with nothing to install - -The [Code Editor demo](../CODE_EDITOR.md) runs a trained water-mapping U-Net on Sentinel-1 -radar imagery. You paste two scripts into the Earth Engine Code Editor, click **Run**, and -export one 224 × 224 pixel prediction to your own project. It needs an Earth Engine account and -a Cloud project [registered for Earth Engine](https://developers.google.com/earth-engine/guides/access). -One export is billed to your project; the first recorded run took 140 s and 64.9 batch -EECU-seconds ([record](evidence/l06d-user-check.json)). - -### 2. Use it from Python - -[Notebook 01, quickstart](../notebooks/01_quickstart.ipynb) -([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/01_quickstart.ipynb)) -does the same thing from a Python notebook: it installs TensorEE, signs in with your own Earth -Engine project, finds a Sentinel-1 scene at a place you pick, runs the same trained U-Net on -one chip there, and shows the radar image and the water map on an interactive map. It uses -only the public weights; you do not need the model's checkpoint file. - -Every notebook first runs offline: it builds the Earth Engine computation without sending it -anywhere, so you can read and run it without an account. The steps that use your Earth Engine -project are switched off until you change `RUN_LIVE = False` to `RUN_LIVE = True` and put your -project ID in place of `YOUR_EE_PROJECT`. Before each of those steps the notebook says what it -will create and what it costs. - -### 3. Bring your own model - -- [Notebook 02, is my model supported?](../notebooks/02_is_my_model_supported.ipynb) - ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)) - needs no Earth Engine account at all. It loads a PyTorch model, checks it, and explains in - plain words what is supported, what is refused and why, and what you can do next. It shows a - model that passes and models that are refused, and has a cell where you put your own. -- [Notebook 03, bring your own CNN](../notebooks/03_bring_your_own_cnn.ipynb) - ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)) - follows a trained segmentation-models-pytorch U-Net from its checkpoint file to Earth - Engine: check it, plan how to cut a region into tiles, upload its weights once, predict one - tile, compare that tile with PyTorch, and export a region. Each step that is billed to your - project is switched off until you turn it on. - -When you want the details, the [Getting started](getting-started.md) page and the -[guides](README.md) are the reference. +1. **See it work, nothing to install.** The [Code Editor demo](../CODE_EDITOR.md): paste two + scripts into the Earth Engine Code Editor and run a trained water-mapping U-Net on one + Sentinel-1 chip. +2. **Use it from Python.** [Notebook 01, quickstart](../notebooks/01_quickstart.ipynb) + ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/01_quickstart.ipynb)): + enter your project, **Run all**, and get a water map at a place you pick. No model file + needed. +3. **Bring your own model.** + - [Notebook 02, is my model supported?](../notebooks/02_is_my_model_supported.ipynb) + ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)): + offline, no account. Checks a PyTorch model and explains what is supported, what is + refused and why. + - [Notebook 03, bring your own CNN](../notebooks/03_bring_your_own_cnn.ipynb) + ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)): + a trained segmentation-models-pytorch U-Net from checkpoint to an exported region. It runs + offline first; each Earth Engine step has its own switch. + +For details, see [Getting started](getting-started.md) and the [guides](README.md). + +## What a run costs + +- Runs use your project's Earth Engine quota (EECU time, Earth Engine's unit of computation). +- That is free on a noncommercial project; a commercial project pays its plan's rate. See + [Earth Engine pricing](https://cloud.google.com/earth-engine/pricing). +- The quickstart's export used about 65 EECU-seconds (64.9 recorded, + [record](evidence/l06d-user-check.json)), plus a few small requests. Notebook 03 says where + a cost has not been measured. ## Words you will meet -- **Chip.** One fixed-size square of pixels that a model takes as input, for example - 224 × 224 pixels of Sentinel-1 at 10 m. The trained U-Net was trained on chips of that size. -- **Tile.** One piece of a larger region. A region too big for one run is cut into tiles that - Earth Engine computes one by one, and the results are put back together. -- **Halo.** The extra border of input pixels around a tile that a model needs to compute the - tile's edge correctly, because each output pixel depends on input pixels around it. The halo - is computed but thrown away, so it is a cost: for a standard U-Net with a ResNet-34 encoder - it is 544 pixels on each side, so a 1,024-pixel tile reads a 2,112-pixel square. -- **Weight asset.** An Earth Engine table asset that holds a model's numbers (its weights). - You upload them once, and every prediction reads them from there. The trained U-Net's weight - assets are public, so you can use them without uploading anything. -- **Export.** An Earth Engine batch task (`Export.image.toAsset` or `toDrive`) that computes an - image and saves it as an asset or a file. Whole models run as exports, because one - interactive request with a whole model exceeded Earth Engine's memory limit. An export runs - in the background: you start it, wait for **Completed** in the Tasks list, then look at the - result. -- **EECU cost.** Earth Engine measures computation in EECU-seconds (Earth Engine Compute - Units), and bills commercial projects for them; see Google's - [Earth Engine pricing](https://cloud.google.com/earth-engine/pricing). Batch exports and - interactive requests are counted separately. The notebooks give the recorded cost where one exists, and say so when a cost has not - been measured. +- **Chip:** the fixed-size square a model takes as input, for example 224 × 224 pixels of + Sentinel-1 at 10 m. +- **Tile:** one piece of a region too large for one run. Tiles are computed one by one and put + back together. +- **Halo:** the extra border of input a tile needs so that its edges come out as the model + would compute them on the whole image. It is read and then dropped: for a U-Net with a + ResNet-34 encoder, 544 pixels per side, so a 1,024-pixel tile reads 2,112 × 2,112. +- **Weight asset:** an Earth Engine table asset holding a model's weights. Uploaded once, + read by every prediction. The trained U-Net's are public. +- **Export:** an Earth Engine batch task that computes an image and saves it as an asset or a + file. Whole models run as exports, because one interactive request with a whole model + exceeded Earth Engine's memory limit. +- **EECU-seconds:** how Earth Engine measures computation time; batch exports and interactive + requests are counted separately. ## What has been checked -TensorEE keeps four kinds of claim apart, and the notebooks do too: - -- **Graph identity**: the Earth Engine computation is exactly, byte for byte, one that was - checked before. -- **Offline agreement**: each building block, computed in NumPy, matches PyTorch. -- **Live agreement**: Earth Engine's result matches PyTorch on the same input, within an - absolute difference of 1e-4 in the logits. For the trained U-Net the recorded single-chip - run differed by at most 1.35e-5 ([validation](validation.md)). -- **Accuracy**: how often the predictions are right. This is not measured by TensorEE. You - measure it on your own labelled data. +- **Graph identity:** the Earth Engine computation is byte for byte one checked before. +- **Offline agreement:** each building block, computed in NumPy, matches PyTorch. +- **Live agreement:** Earth Engine's result matches PyTorch on the same input, within 1e-4 in + the logits. For the trained U-Net's recorded chip: at most 1.35e-5 + ([validation](validation.md)). +- **Accuracy:** not measured by TensorEE; measure it on your own labelled data. -If a notebook gets a result you did not expect, or refuses a model you think should work, -[open an issue](https://github.com/noelEOS/tensor-ee/issues) with the message it printed. +Something unexpected, or a model refused that you think should work? +[Open an issue](https://github.com/noelEOS/tensor-ee/issues) with the message it printed. diff --git a/notebooks/01_quickstart.ipynb b/notebooks/01_quickstart.ipynb index 330418e..7418b40 100644 --- a/notebooks/01_quickstart.ipynb +++ b/notebooks/01_quickstart.ipynb @@ -5,597 +5,250 @@ "id": "2ac26326", "metadata": {}, "source": [ - "# TensorEE quickstart: a trained water U-Net on one Sentinel-1 chip\n", + "# TensorEE quickstart: map water with a trained U-Net, inside Earth Engine\n", "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/01_quickstart.ipynb)\n", "\n", - "This notebook runs a trained neural network inside Google Earth Engine. The model is a U-Net\n", - "with a ResNet-18 encoder that maps open water in Sentinel-1 radar imagery. It takes one chip of\n", - "224 × 224 pixels at 10 m (2.24 km on a side) and returns, for every pixel, the probability\n", - "that it is water and a water mask.\n", - "\n", - "Earth Engine computes the whole prediction from your Sentinel-1 scene. Nothing is downloaded\n", - "and no model server is involved: the network's arithmetic is a large Earth Engine computation,\n", - "and its weights are public Earth Engine table assets. You do not need the model's checkpoint\n", - "file.\n", - "\n", - "**What you need.** Nothing, to read and run the notebook offline. To run the model, an Earth\n", - "Engine account and a Cloud project\n", - "[registered for Earth Engine](https://developers.google.com/earth-engine/guides/access) where\n", - "you can create an image asset.\n", - "\n", - "**What it costs.** Offline, nothing. With `RUN_LIVE = True` the notebook makes a few small\n", - "requests (to place the chip and find a scene) and **one batch export of one chip**, billed to\n", - "your project. The first recorded export of this chip took 140 s and 64.9 batch EECU-seconds.\n", - "The export creates one image asset in your project (about 1 MB). Each step says what it does\n", - "before it does it." + "- **What it does:** runs a trained water-mapping U-Net on one Sentinel-1 chip\n", + " (224 × 224 pixels at 10 m) at a place you pick, and shows the water map.\n", + "- **How:** Earth Engine computes the whole network itself. Nothing is downloaded, and the\n", + " weights are public Earth Engine assets, so you need no model file.\n", + "- **You need:** an Earth Engine account and a Cloud project\n", + " [registered for Earth Engine](https://developers.google.com/earth-engine/guides/access).\n", + "- **To run:** enter your project in step 2, then **Runtime → Run all**.\n", + "- **Cost:** the run uses your project's Earth Engine quota (EECU time). That is free on a\n", + " noncommercial project; a commercial project pays its plan's rate. The example's export used\n", + " about 65 EECU-seconds (64.9 recorded), plus a few small requests." ] }, { - "cell_type": "markdown", + "cell_type": "code", "id": "7361375f", "metadata": {}, "source": [ - "## 1. Install\n", - "\n", - "The trained U-Net's Earth Engine computation is stored in the repository, in the file the\n", - "[Code Editor demo](https://github.com/noelEOS/tensor-ee/blob/main/CODE_EDITOR.md) uses. That\n", - "file is not part of the package on PyPI, so this cell copies the repository and installs\n", - "TensorEE from it. It also installs `folium` for the map. In Colab it takes a minute or two.\n", + "# @title 1. Install (about a minute)\n", + "# The model's Earth Engine computation is in the repository, not on PyPI, so install from it.\n", + "!git clone --quiet --depth 1 https://github.com/noelEOS/tensor-ee.git\n", + "%pip install --quiet ./tensor-ee folium\n", "\n", - "If you opened this notebook from your own copy of the repository, you can skip this cell." - ] - }, - { - "cell_type": "code", - "id": "33d242a2", - "metadata": {}, - "source": [ - "# The Code Editor module (demos/code_editor/segwater.js) is in the repository, not on PyPI,\n", - "# so this notebook installs TensorEE from a copy of the repository.\n", - "!git clone --depth 1 https://github.com/noelEOS/tensor-ee.git\n", - "%pip install --quiet ./tensor-ee folium" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "9b2cdbb8", - "metadata": {}, - "source": [ "import sys\n", "from pathlib import Path\n", "\n", - "\n", - "def find_repository():\n", - " \"\"\"The copy of the repository: ./tensor-ee in Colab, or the folder this notebook is in.\"\"\"\n", - " for folder in (Path(\"tensor-ee\"), Path(\".\"), Path(\"..\")):\n", - " if (folder / \"demos\" / \"code_editor\" / \"segwater.js\").is_file():\n", - " return folder.resolve()\n", - " raise FileNotFoundError(\"Run the install cell first: it copies the repository to ./tensor-ee\")\n", - "\n", - "\n", - "REPOSITORY = find_repository()\n", - "sys.path.insert(0, str(REPOSITORY))\n", - "print(\"Using the repository in\", REPOSITORY)" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "376fa9ae", - "metadata": {}, - "source": [ - "## 2. Settings\n", - "\n", - "This is the only cell you need to edit.\n", - "\n", - "- `PROJECT`: your Cloud project ID, the one you use in the Code Editor.\n", - "- `RUN_LIVE`: leave it `False` to run everything offline first. Set it to `True` to sign in\n", - " and run the steps that use your project.\n", - "- `LONGITUDE`, `LATITUDE`: the centre of the chip, in degrees. The default is Sun Moon Lake,\n", - " Taiwan, where the model's Earth Engine result was checked against PyTorch.\n", - "- `START_DATE`, `END_DATE`: the period in which to look for a Sentinel-1 scene.\n", - "- `ASSET_NAME`: the name of the new image asset the export creates in your project. Use a new\n", - " name for each run: an export never overwrites an asset.\n", - "\n", - "After changing `RUN_LIVE`, restart the runtime (in Colab: **Runtime → Restart session**) and\n", - "run the notebook again from the top: the offline mode stays installed until the runtime\n", - "restarts." - ] - }, - { - "cell_type": "code", - "id": "c035f536", - "metadata": {}, - "source": [ - "PROJECT = \"YOUR_EE_PROJECT\" # your Cloud project registered for Earth Engine\n", - "RUN_LIVE = False # True: sign in and run the steps that use your project\n", - "\n", - "LONGITUDE, LATITUDE = 120.9150, 23.8620 # Sun Moon Lake, Taiwan\n", - "START_DATE, END_DATE = \"2024-01-01\", \"2024-01-08\" # look for Sentinel-1 scenes in this period\n", - "ASSET_NAME = \"tensor_ee_quickstart_chip\" # a new image asset in your project" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "5b6ecfc4", - "metadata": {}, - "source": [ - "PLACEHOLDER = \"YOUR_EE_PROJECT\"\n", - "\n", - "\n", - "def live_step(what):\n", - " \"\"\"True when a step that uses Earth Engine should run; it refuses a placeholder project.\"\"\"\n", - " if not RUN_LIVE:\n", - " print(f\"Skipped: {what}. RUN_LIVE is False, so nothing was sent to Earth Engine.\")\n", - " return False\n", - " if not PROJECT or PROJECT == PLACEHOLDER:\n", - " raise ValueError(\"Set PROJECT to your own Cloud project ID in the settings cell \"\n", - " \"before running with RUN_LIVE = True.\")\n", - " return True" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "d32406d8", - "metadata": {}, - "source": [ - "## 3. Sign in, or work offline\n", - "\n", - "With `RUN_LIVE = True` the next cell signs you in. The first time, `ee.Authenticate()` shows\n", - "a link or a prompt; follow it with the Google account you use for Earth Engine. Every request\n", - "after that is made with your account and billed to `PROJECT`.\n", - "\n", - "With `RUN_LIVE = False`, the cell after it installs the Earth Engine client's offline\n", - "function signatures (`ee.apitestcase.ApiTestCase`, a test helper that ships with the client).\n", - "They let the notebook build the same Earth Engine computation without credentials and without\n", - "sending anything." - ] - }, - { - "cell_type": "code", - "id": "ee640c02", - "metadata": { - "tags": [ - "live" - ] - }, - "source": [ - "import ee\n", - "\n", - "if live_step(\"sign in to Earth Engine\"):\n", - " ee.Authenticate()\n", - " ee.Initialize(project=PROJECT)\n", - " print(\"Signed in. Requests are billed to\", PROJECT)" + "REPOSITORY = next(folder.resolve() for folder in (Path(\"tensor-ee\"), Path(\".\"), Path(\"..\"))\n", + " if (folder / \"demos\" / \"code_editor\" / \"segwater.js\").is_file())\n", + "sys.path.insert(0, str(REPOSITORY))" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", - "id": "d1822a6b", - "metadata": {}, - "source": [ - "import ee\n", - "\n", - "if not RUN_LIVE:\n", - " from ee.apitestcase import ApiTestCase\n", - "\n", - " offline = ApiTestCase()\n", - " offline.setUp() # offline signatures: no credentials, no network\n", - " print(\"Offline: the notebook builds Earth Engine computations but sends none.\")" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "6b906da1", - "metadata": {}, - "source": [ - "## 4. Load the trained model\n", - "\n", - "The model comes as one file, `demos/code_editor/segwater.js`. TensorEE built it from the\n", - "trained checkpoint: it holds the complete Earth Engine computation for one chip, with blanks\n", - "for the place and the scene, and it reads the network's weights from 29 public Earth Engine\n", - "table assets. That is why you need neither the checkpoint nor an upload. The same file runs in\n", - "the Code Editor demo.\n", - "\n", - "The file stores the computation compressed. Reading it checks that it is complete and\n", - "undamaged." - ] - }, - { - "cell_type": "code", - "id": "e68ae525", - "metadata": {}, - "source": [ - "from demos.code_editor import build as segwater_module\n", - "\n", - "TEMPLATE, METADATA = segwater_module.parse_module(\n", - " (REPOSITORY / \"demos\" / \"code_editor\" / \"segwater.js\").read_text())\n", - "\n", - "print(\"Model computation:\", f\"{METADATA['templateBytes']:,}\", \"characters, checked intact\")\n", - "print(\"Weight tables it reads:\", len(METADATA[\"weightAssets\"]))\n", - "print(\"Built from the checkpoint with SHA-256\", METADATA[\"checkpointSha256\"][:16] + \"...\")" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "52b594cb", + "id": "33d242a2", "metadata": {}, "source": [ - "## 5. Choose the place\n", + "# @title 2. Your project and place\n", + "# @markdown - **PROJECT**: your Cloud project ID, as in the Code Editor. Nothing else needs editing.\n", + "# @markdown - The default place is Sun Moon Lake, Taiwan. The date range picks a Sentinel-1 scene.\n", + "PROJECT = \"\" # @param {type:\"string\"}\n", + "LONGITUDE = 120.915 # @param {type:\"number\"}\n", + "LATITUDE = 23.862 # @param {type:\"number\"}\n", + "START_DATE = \"2024-01-01\" # @param {type:\"date\"}\n", + "END_DATE = \"2024-01-08\" # @param {type:\"date\"}\n", "\n", - "The chip is a 224 × 224 grid of 10 m pixels in the local UTM zone, centred on your point.\n", - "The model was trained on chips of exactly this size, so the notebook keeps it.\n", + "import time\n", "\n", - "Offline, the notebook uses the preset chip at Sun Moon Lake. With `RUN_LIVE = True` the live\n", - "cell asks Earth Engine to convert your longitude and latitude to UTM metres: **one small\n", - "request**." - ] - }, - { - "cell_type": "code", - "id": "be96d487", - "metadata": {}, - "source": [ - "from tensor_ee import Grid\n", + "ASSET_NAME = \"tensor_ee_quickstart_\" + time.strftime(\"%Y%m%d_%H%M%S\") # a new asset each run\n", "\n", - "CHIP, PIXEL = 224, 10 # pixels per side, metres per pixel\n", - "PRESET = { # the verified Sun Moon Lake chip of the Code Editor demo\n", - " \"crs\": \"EPSG:32651\", \"x\": 286030, \"y\": 2640860,\n", - " \"scene\": \"COPERNICUS/S1_GRD/S1A_IW_GRDH_1SDV_20240103T100134_20240103T100159_051941_06469E_B0BD\",\n", - "}\n", "\n", + "class Stop(Exception):\n", + " \"\"\"Stops the notebook with a one-line message instead of a traceback.\"\"\"\n", "\n", - "def utm_zone(longitude, latitude):\n", - " \"\"\"The EPSG code of the WGS 84 UTM zone containing a point.\"\"\"\n", - " zone = min(60, max(1, int((longitude + 180) // 6) + 1))\n", - " return f\"EPSG:{326 if latitude >= 0 else 327}{zone:02d}\"\n", + " def _render_traceback_(self):\n", + " return [f\"Stopped: {self}\"]\n", "\n", "\n", - "place = {\"crs\": PRESET[\"crs\"], \"x\": PRESET[\"x\"], \"y\": PRESET[\"y\"]} # offline: the preset" + "if not PROJECT.strip():\n", + " raise Stop(\"enter your Cloud project ID in PROJECT above, then run all cells again.\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", - "id": "33726682", + "id": "9b2cdbb8", "metadata": { "tags": [ "live" ] }, "source": [ - "if live_step(\"place the chip at your point (1 request)\"):\n", - " crs = utm_zone(LONGITUDE, LATITUDE)\n", - " x, y = ee.Geometry.Point([LONGITUDE, LATITUDE]).transform(crs, 0.01).coordinates().getInfo()\n", - " half = CHIP * PIXEL // 2\n", - " # The chip's upper-left corner, on whole 10 m steps.\n", - " place = {\"crs\": crs, \"x\": round((x - half) / PIXEL) * PIXEL,\n", - " \"y\": round((y + half) / PIXEL) * PIXEL}" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "61805fcd", - "metadata": {}, - "source": [ - "grid = Grid(place[\"crs\"], (PIXEL, 0, place[\"x\"], 0, -PIXEL, place[\"y\"]), CHIP, CHIP)\n", - "print(\"Chip grid:\", grid.crs, \"upper-left corner\", (place[\"x\"], place[\"y\"]),\n", - " f\"{CHIP} x {CHIP} pixels of {PIXEL} m\")" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "3a027ecd", - "metadata": {}, - "source": [ - "## 6. Find a Sentinel-1 scene\n", - "\n", - "The model reads the VV and VH bands of a Sentinel-1 GRD scene (interferometric wide swath, 10 m)\n", - "in decibels, as Earth Engine stores them. The scene must cover the whole chip.\n", + "# @title 3. Sign in\n", + "# @markdown The first time, follow the link with the Google account you use for Earth Engine.\n", + "import ee\n", "\n", - "With `RUN_LIVE = True` the live cell lists up to ten scenes in your period and, for each, the\n", - "fraction of the chip it covers, and picks the first that covers it all: **one request**, which\n", - "reads only the chip's area. Offline, the notebook uses the preset scene of 3 January 2024." - ] - }, - { - "cell_type": "code", - "id": "05f143cc", - "metadata": {}, - "source": [ - "scene_id = PRESET[\"scene\"] # offline: the preset scene" + "ee.Authenticate()\n", + "ee.Initialize(project=PROJECT)\n", + "print(\"Signed in to\", PROJECT)" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", - "id": "4d94e7c5", + "id": "376fa9ae", "metadata": { "tags": [ "live" ] }, "source": [ - "if live_step(\"find a Sentinel-1 scene (1 request)\"):\n", - " area = grid.geometry()\n", - " collection = (ee.ImageCollection(\"COPERNICUS/S1_GRD\")\n", - " .filterBounds(area).filterDate(START_DATE, END_DATE)\n", - " .filter(ee.Filter.eq(\"instrumentMode\", \"IW\"))\n", - " .filter(ee.Filter.listContains(\"transmitterReceiverPolarisation\", \"VV\"))\n", - " .filter(ee.Filter.listContains(\"transmitterReceiverPolarisation\", \"VH\"))\n", - " .filter(ee.Filter.eq(\"resolution_meters\", 10))\n", - " .sort(\"system:time_start\"))\n", - "\n", - " def describe(image):\n", - " image = ee.Image(image)\n", - " valid = (image.select([\"VV\", \"VH\"]).mask().reduce(ee.Reducer.min()).gt(0)\n", - " .unmask(0, False).rename(\"valid\"))\n", - " coverage = valid.reduceRegion(reducer=ee.Reducer.mean(), geometry=area, crs=grid.crs,\n", - " crsTransform=list(grid.transform),\n", - " maxPixels=CHIP * CHIP, tileScale=2).get(\"valid\")\n", - " return ee.Dictionary({\n", - " \"id\": ee.String(\"COPERNICUS/S1_GRD/\").cat(image.get(\"system:index\")),\n", - " \"time\": image.date().format(\"YYYY-MM-dd HH:mm\"),\n", - " \"pass\": image.get(\"orbitProperties_pass\"), \"coverage\": coverage})\n", + "# @title 4. Find the chip and a Sentinel-1 scene (2 small requests)\n", + "# @markdown A 224 × 224 chip of 10 m pixels in the local UTM zone, centred on your point,\n", + "# @markdown and the first scene in the date range that covers all of it.\n", + "from tensor_ee import Grid\n", "\n", - " scenes = collection.toList(10).map(describe).getInfo()\n", - " for scene in scenes:\n", - " print(f\"{scene['time']} {scene['pass']:<10} covers {100 * (scene['coverage'] or 0):6.2f}% {scene['id']}\")\n", - " complete = [scene for scene in scenes if (scene[\"coverage\"] or 0) >= 0.999]\n", - " if not complete:\n", - " raise RuntimeError(\"No scene in this period covers the whole chip. Widen START_DATE to \"\n", - " \"END_DATE, or move the point away from the edge of the swath.\")\n", - " scene_id = complete[0][\"id\"]\n", - " print(\"Using\", scene_id)" + "CHIP, PIXEL = 224, 10\n", + "zone = min(60, max(1, int((LONGITUDE + 180) // 6) + 1))\n", + "crs = f\"EPSG:{326 if LATITUDE >= 0 else 327}{zone:02d}\"\n", + "x, y = ee.Geometry.Point([LONGITUDE, LATITUDE]).transform(crs, 0.01).coordinates().getInfo()\n", + "half = CHIP * PIXEL // 2\n", + "place = {\"crs\": crs, \"x\": round((x - half) / PIXEL) * PIXEL, \"y\": round((y + half) / PIXEL) * PIXEL}\n", + "grid = Grid(crs, (PIXEL, 0, place[\"x\"], 0, -PIXEL, place[\"y\"]), CHIP, CHIP)\n", + "\n", + "scenes = (ee.ImageCollection(\"COPERNICUS/S1_GRD\")\n", + " .filterBounds(grid.geometry()).filterDate(START_DATE, END_DATE)\n", + " .filter(ee.Filter.eq(\"instrumentMode\", \"IW\"))\n", + " .filter(ee.Filter.listContains(\"transmitterReceiverPolarisation\", \"VV\"))\n", + " .filter(ee.Filter.listContains(\"transmitterReceiverPolarisation\", \"VH\"))\n", + " .filter(ee.Filter.eq(\"resolution_meters\", 10))\n", + " .sort(\"system:time_start\").limit(10))\n", + "\n", + "\n", + "def coverage(image):\n", + " valid = image.select([\"VV\", \"VH\"]).mask().reduce(ee.Reducer.min()).gt(0).unmask(0, False)\n", + " share = valid.rename(\"valid\").reduceRegion(\n", + " reducer=ee.Reducer.mean(), geometry=grid.geometry(), crs=grid.crs,\n", + " crsTransform=list(grid.transform), maxPixels=CHIP * CHIP, tileScale=2).get(\"valid\")\n", + " return ee.Feature(None, {\"id\": ee.String(\"COPERNICUS/S1_GRD/\").cat(image.get(\"system:index\")),\n", + " \"share\": share})\n", + "\n", + "\n", + "found = scenes.map(coverage).reduceColumns(ee.Reducer.toList(2), [\"id\", \"share\"]).get(\"list\").getInfo()\n", + "complete = [scene for scene, share in found if share >= 0.999]\n", + "if not complete:\n", + " raise Stop(f\"no Sentinel-1 scene covers the whole chip between {START_DATE} and {END_DATE}; \"\n", + " \"widen the dates or move the point.\")\n", + "scene_id = complete[0]\n", + "print(\"Chip:\", crs, (place[\"x\"], place[\"y\"]), \"| Scene:\", scene_id)" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "ef2053c6", - "metadata": {}, - "source": [ - "## 7. Build the prediction\n", - "\n", - "This cell fills the model's blanks with your grid and scene and turns the text back into an\n", - "Earth Engine image. Nothing is computed yet: like any Earth Engine image, it is a recipe that\n", - "Earth Engine evaluates only when you export or display it.\n", - "\n", - "The recipe converts the radar values from decibels the way the model was trained, runs the\n", - "U-Net, and returns five bands:\n", - "\n", - "| Band | Meaning |\n", - "|---|---|\n", - "| `logit_0`, `logit_1` | the network's raw scores for \"not water\" and \"water\" |\n", - "| `water_probability` | the probability of water, from 0 to 1 |\n", - "| `water_mask` | 1 where the probability is at least 0.5, else 0 |\n", - "| `coverage` | 1 where both radar bands and the result are valid |\n", - "\n", - "For the preset chip, the notebook also checks the computation against the record: it should\n", - "be identical, byte for byte, to the one TensorEE's own Python recipe builds with the trained\n", - "checkpoint, which the repository records." - ] - }, { "cell_type": "code", - "id": "6a971198", + "id": "c035f536", "metadata": {}, "source": [ - "import json\n", + "# @title 5. Build the model's computation\n", + "# @markdown Fills the trained U-Net's stored computation with your chip and scene. Nothing runs yet.\n", + "import ee\n", + "from demos.code_editor import build as segwater_module\n", + "from tensor_ee import Grid\n", "\n", - "sys.setrecursionlimit(max(sys.getrecursionlimit(), 20000)) # the U-Net is a deep graph\n", - "OUTPUT_BANDS = list(segwater_module.OUTPUT_BANDS)\n", + "sys.setrecursionlimit(max(sys.getrecursionlimit(), 20000)) # the network is a deep expression\n", + "TEMPLATE, METADATA = segwater_module.parse_module(\n", + " (REPOSITORY / \"demos\" / \"code_editor\" / \"segwater.js\").read_text())\n", + "BANDS = list(segwater_module.OUTPUT_BANDS)\n", "\n", + "grid = Grid(place[\"crs\"], (10, 0, place[\"x\"], 0, -10, place[\"y\"]), 224, 224)\n", "text = segwater_module.substitute(TEMPLATE, grid.crs, scene_id,\n", " (place[\"x\"], place[\"y\"]), (place[\"x\"], place[\"y\"]))\n", "chip = ee.Image(ee.deserializer.fromJSON(text))\n", - "# Put the chip on its grid, as the library's recipe does for every chip it computes.\n", - "prediction = (ee.ImageCollection.fromImages([chip]).mosaic().select(OUTPUT_BANDS)\n", + "prediction = (ee.ImageCollection.fromImages([chip]).mosaic().select(BANDS)\n", " .setDefaultProjection(grid.projection()).clip(grid.geometry())\n", " .updateMask(grid.mask()).toFloat())\n", - "\n", - "serialized = prediction.serialize()\n", - "graph_sha256 = segwater_module.sha256(serialized)\n", - "print(\"Bands:\", OUTPUT_BANDS)\n", - "print(f\"The computation is {len(serialized):,} characters of Earth Engine expression.\")\n", - "\n", - "record = json.loads((REPOSITORY / \"docs\" / \"evidence\" / \"s03-template-identity.json\").read_text())\n", - "preset = next(aoi for aoi in record[\"aois\"] if aoi[\"name\"] == \"segwater_sun_moon_lake_1x1\")\n", - "if (scene_id, grid.crs, place[\"x\"], place[\"y\"]) != (PRESET[\"scene\"], PRESET[\"crs\"], PRESET[\"x\"], PRESET[\"y\"]):\n", - " print(\"Your own place: there is no recorded computation to compare it with.\")\n", - "elif ee.__version__ != METADATA[\"earthengineApiVersion\"]:\n", - " print(f\"Not compared: your Earth Engine client is {ee.__version__}; the record was made with \"\n", - " f\"{METADATA['earthengineApiVersion']}, and client versions write expressions differently.\")\n", - "elif graph_sha256 == preset[\"library_sha256\"]:\n", - " print(\"Identical, byte for byte, to the recorded computation of the library's recipe.\")\n", - "else:\n", - " print(\"Different from the recorded computation: please report this as an issue.\")" + "print(\"Output bands:\", \", \".join(BANDS))" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "21a05245", - "metadata": {}, - "source": [ - "## 8. Export the prediction (billed to your project)\n", - "\n", - "**What this creates.** One batch export task, `Export.image.toAsset`, and when it completes,\n", - "one new image asset, `projects/PROJECT/assets/ASSET_NAME`: five float32 bands over\n", - "224 × 224 pixels, about 1 MB.\n", - "\n", - "**What it costs.** The first recorded export of the Sun Moon Lake chip took 140 s and\n", - "64.9 batch EECU-seconds. No cost has been recorded for other places. How long an export\n", - "waits and runs varies with Earth Engine's load.\n", - "\n", - "The model runs as an export, not as map tiles: when a whole trained U-Net was sent as one\n", - "interactive request, it exceeded Earth Engine's memory limit. The cell after this one waits for the task." - ] - }, { "cell_type": "code", - "id": "ce479c84", + "id": "5b6ecfc4", "metadata": { "tags": [ "live" ] }, "source": [ + "# @title 6. Run the model (an export of about 65 EECU-seconds, a few minutes)\n", + "# @markdown Earth Engine computes the chip in a batch export and saves it as a new image asset.\n", "asset_id = f\"projects/{PROJECT}/assets/{ASSET_NAME}\"\n", - "\n", - "if live_step(\"export the prediction (1 batch task, 1 new asset)\"):\n", - " payload = prediction.set({\"source_scene\": scene_id,\n", - " \"checkpoint_sha256\": METADATA[\"checkpointSha256\"],\n", - " \"template_sha256\": METADATA[\"templateSha256\"]})\n", - " task = ee.batch.Export.image.toAsset(\n", - " image=payload, description=\"tensor_ee_quickstart\", assetId=asset_id,\n", - " region=grid.geometry(), crs=grid.crs, crsTransform=list(grid.transform),\n", - " maxPixels=CHIP * CHIP * len(OUTPUT_BANDS), pyramidingPolicy={\".default\": \"sample\"})\n", - " task.start()\n", - " print(\"Started task\", task.id, \"writing\", asset_id)" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "2d79e129", - "metadata": {}, - "source": [ - "Waiting checks the task's status every 20 seconds. You can also follow it in the **Tasks** tab\n", - "of the Code Editor. If you stop waiting, the export carries on; run this cell again later." - ] - }, - { - "cell_type": "code", - "id": "7e651659", - "metadata": { - "tags": [ - "live" - ] - }, - "source": [ - "import time\n", - "\n", - "if live_step(\"wait for the export\"):\n", - " deadline = time.time() + 30 * 60\n", - " while True:\n", - " state = task.status()[\"state\"]\n", - " if state in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", - " break\n", - " if time.time() > deadline:\n", - " raise TimeoutError(\"Still running after 30 minutes; run this cell again later.\")\n", - " time.sleep(20)\n", - " print(\"Export\", state)\n", - " if state != \"COMPLETED\":\n", - " raise RuntimeError(task.status().get(\"error_message\", \"the export did not complete\"))" + "task = ee.batch.Export.image.toAsset(\n", + " image=prediction.set({\"source_scene\": scene_id}), description=\"tensor_ee_quickstart\",\n", + " assetId=asset_id, region=grid.geometry(), crs=grid.crs, crsTransform=list(grid.transform),\n", + " maxPixels=224 * 224 * len(BANDS), pyramidingPolicy={\".default\": \"sample\"})\n", + "task.start()\n", + "print(\"Running; you can also watch it in the Code Editor's Tasks tab.\")\n", + "while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", + " time.sleep(20)\n", + "if state != \"COMPLETED\":\n", + " raise Stop(f\"the export ended {state}: {task.status().get('error_message', '')}\")\n", + "print(\"Done:\", asset_id)" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "92dbf3c4", - "metadata": {}, - "source": [ - "## 9. See the radar image and the water map\n", - "\n", - "The map shows the Sentinel-1 VV band and the completed asset's water probability and mask.\n", - "The layers are Earth Engine map tiles: displaying them reads the saved asset and the scene,\n", - "so panning does not run the network again. Each tile you view is a small interactive request.\n", - "Use the layer control at the top right to switch layers." - ] - }, { "cell_type": "code", - "id": "f67c0038", + "id": "d32406d8", "metadata": { "tags": [ "live" ] }, "source": [ - "if live_step(\"show the result on a map\"):\n", - " import folium\n", - " from IPython.display import display\n", - "\n", - " def tile_layer(image, visualization, name, show=True):\n", - " url = ee.Image(image).getMapId(visualization)[\"tile_fetcher\"].url_format\n", - " return folium.raster_layers.TileLayer(\n", - " tiles=url, name=name, overlay=True, show=show,\n", - " attr=\"Google Earth Engine; contains modified Copernicus Sentinel data\")\n", - "\n", - " result = ee.Image(asset_id)\n", - " outline = grid.geometry().transform(\"EPSG:4326\", 0.01).getInfo()\n", - " longitudes, latitudes = zip(*outline[\"coordinates\"][0])\n", - " water_map = folium.Map(location=[sum(latitudes) / len(latitudes), sum(longitudes) / len(longitudes)],\n", - " zoom_start=14)\n", - " tile_layer(ee.Image(scene_id).select(\"VV\"), {\"min\": -25, \"max\": 0},\n", - " \"Sentinel-1 VV (dB)\").add_to(water_map)\n", - " tile_layer(result.select(\"water_probability\"),\n", - " {\"min\": 0, \"max\": 1, \"palette\": [\"ffffff\", \"08306b\"]},\n", - " \"Water probability\").add_to(water_map)\n", - " tile_layer(result.select(\"water_mask\").selfMask(), {\"palette\": [\"00bfff\"]},\n", - " \"Water mask\", show=False).add_to(water_map)\n", - " folium.GeoJson(outline, name=\"Chip\", style_function=lambda _: {\n", - " \"fill\": False, \"color\": \"#ff7f00\", \"weight\": 2}).add_to(water_map)\n", - " folium.LayerControl().add_to(water_map)\n", - " display(water_map)" + "# @title 7. See the radar image and the water map\n", + "# @markdown Use the layer control (top right) to switch layers.\n", + "import folium\n", + "\n", + "def layer(image, visualization, name, show=True):\n", + " url = ee.Image(image).getMapId(visualization)[\"tile_fetcher\"].url_format\n", + " return folium.raster_layers.TileLayer(\n", + " tiles=url, name=name, overlay=True, show=show,\n", + " attr=\"Google Earth Engine; contains modified Copernicus Sentinel data\")\n", + "\n", + "result = ee.Image(asset_id)\n", + "water_map = folium.Map(location=[LATITUDE, LONGITUDE], zoom_start=14)\n", + "layer(ee.Image(scene_id).select(\"VV\"), {\"min\": -25, \"max\": 0}, \"Sentinel-1 VV (dB)\").add_to(water_map)\n", + "layer(result.select(\"water_probability\"), {\"min\": 0, \"max\": 1, \"palette\": [\"ffffff\", \"08306b\"]},\n", + " \"Water probability\").add_to(water_map)\n", + "layer(result.select(\"water_mask\").selfMask(), {\"palette\": [\"00bfff\"]}, \"Water mask\",\n", + " show=False).add_to(water_map)\n", + "folium.LayerControl().add_to(water_map)\n", + "water_map" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", - "id": "582437dd", + "id": "ee640c02", "metadata": {}, "source": [ - "## What this shows, and what it does not\n", + "## Good to know\n", "\n", - "- **The computation.** For the preset chip, the notebook's computation is identical to the\n", - " library's own recipe built with the trained checkpoint (step 7). For another place it is the\n", - " same computation with a different grid and scene.\n", - "- **Earth Engine against PyTorch.** For the preset chip, an earlier export was compared pixel by\n", - " pixel with the trained model in PyTorch: the largest logit difference was 1.35e-5, the largest\n", - " probability difference 6.4e-6, and the water masks agreed on all 50,176 pixels\n", - " ([validation](https://github.com/noelEOS/tensor-ee/blob/main/docs/validation.md)). Your\n", - " export at another place has not been compared; that needs the checkpoint, which is not yet\n", - " public.\n", - "- **Accuracy.** The water map is a prediction, not ground truth, and this notebook does not\n", - " measure how often it is right. The model was trained on Sentinel-1 chips with water labels\n", - " derived from Sentinel-2\n", + "- **Checked, not guaranteed accurate.** At Sun Moon Lake, this computation matched the trained\n", + " model in PyTorch to 1.35e-5 in the logits, with identical water masks\n", + " ([validation](https://github.com/noelEOS/tensor-ee/blob/main/docs/validation.md)). That is\n", + " agreement, not accuracy: the map is a prediction, not ground truth\n", " ([model notes](https://github.com/noelEOS/tensor-ee/blob/main/docs/model.md)).\n", - "- **Size.** One chip of 224 × 224 pixels per export. The Code Editor demo runs up to four chips\n", - " side by side; each chip is computed on its own, so seams can appear at chip borders.\n", - "\n", - "**Cleaning up.** The asset stays in your project until you delete it, in the Code Editor's\n", - "**Assets** tab.\n", - "\n", - "**Next.** [Notebook 02](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)\n", - "checks whether your own model can run this way, and\n", - "[notebook 03](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", - "takes a trained U-Net of your own to Earth Engine." + "- **Why an export:** the whole network in one interactive request exceeds Earth Engine's memory\n", + " limit, so it runs as a batch task. The map then reads the saved asset, which stays in your\n", + " project until you delete it (Code Editor, **Assets** tab).\n", + "- **Next:** [is my model supported?](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)\n", + " and [bring your own CNN](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb);\n", + " [start here](https://github.com/noelEOS/tensor-ee/blob/main/docs/start-here.md) for the background." ] } ], diff --git a/notebooks/02_is_my_model_supported.ipynb b/notebooks/02_is_my_model_supported.ipynb index 6ccaadf..d6ee1fd 100644 --- a/notebooks/02_is_my_model_supported.ipynb +++ b/notebooks/02_is_my_model_supported.ipynb @@ -9,22 +9,15 @@ "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)\n", "\n", - "TensorEE can run a model in Earth Engine only if every operation in it has a TensorEE building\n", - "block, at settings the library can compute exactly. This notebook checks that for a PyTorch\n", - "model, and explains the answer in plain words: what is supported, what is refused and why,\n", - "and what you can do next.\n", - "\n", - "**It runs entirely on your computer or in Colab.** It needs no Earth Engine account, makes no\n", - "Earth Engine request and downloads no weights: every example model is built with random\n", - "weights, which is all a support check needs. The check looks at the model's structure, never\n", - "at its predictions.\n", - "\n", - "The notebook works through four examples, then has a cell for your own model:\n", - "\n", - "1. a U-Net that is supported;\n", - "2. two models that are refused, and what the refusal says;\n", - "3. an attention model, which goes through a different path, the converter;\n", - "4. your model." + "TensorEE runs a model only if every operation in it has a building block it can compute\n", + "exactly. This notebook checks a PyTorch model and says what is supported, what is refused, and\n", + "what to do next.\n", + "\n", + "- **Fully offline:** no Earth Engine account, no requests, no downloads. The examples use\n", + " random weights, which is all a support check needs: it reads the model's structure, not its\n", + " predictions.\n", + "- **To run:** **Runtime → Run all**. Then put your own model in step 6.\n", + "- **Examples:** a supported U-Net; two refused models; an attention model through the converter." ] }, { @@ -34,8 +27,8 @@ "source": [ "## 1. Install\n", "\n", - "The `swin` extra brings segmentation-models-pytorch 0.5.0 and timm 1.0.30, the versions the\n", - "library was checked with. PyTorch is already installed in Colab." + "- The `swin` extra brings segmentation-models-pytorch 0.5.0 and timm 1.0.30, the versions the\n", + " library was checked with." ] }, { @@ -56,10 +49,9 @@ "id": "c424214b", "metadata": {}, "source": [ - "## 2. Settings\n", + "## 2. Setup\n", "\n", - "This notebook has no live step, so it has no project to set: nothing here contacts Earth\n", - "Engine. The only setting is the size of the chip the checks use." + "- The chip size the examples are checked at, and a helper that explains a report." ] }, { @@ -67,7 +59,6 @@ "id": "295d1ada", "metadata": {}, "source": [ - "RUN_LIVE = False # this notebook never contacts Earth Engine\n", "INPUT_SHAPE = (1, 2, 64, 64) # one chip: (1, bands, height, width)" ], "execution_count": null, @@ -126,16 +117,10 @@ "source": [ "## 3. A supported model: a U-Net with a ResNet-18 encoder\n", "\n", - "This is the architecture of the trained water model in notebook 01, built here with random\n", - "weights (`encoder_weights=None` means nothing is downloaded).\n", - "\n", - "Two calls do the work:\n", - "\n", - "- `adapt_smp_model` prepares a segmentation-models-pytorch model for TensorEE at a fixed chip\n", - " size. It also runs the original model and the prepared one on one test input and checks that\n", - " they give exactly the same numbers.\n", - "- `check_compatibility` walks through every operation of the prepared model and says whether\n", - " TensorEE can compute it." + "- The architecture of notebook 01's water model, with random weights.\n", + "- `adapt_smp_model` prepares an SMP model at a fixed chip size, and checks that it gives exactly\n", + " the original's numbers on a test input.\n", + "- `check_compatibility` checks every operation." ] }, { @@ -157,16 +142,13 @@ "id": "443e3ab9", "metadata": {}, "source": [ - "\"Supported\" means the library can express this model, operation by operation, as Earth Engine\n", - "computations. It is not yet a claim about the numbers. Separately, the library has evidence\n", - "for parts of the path: the building blocks agree with PyTorch at the settings they were\n", - "validated at, and one trained U-Net with a ResNet-18 encoder, prepared this way, ran in Earth\n", - "Engine on a 64 × 64 test input and matched PyTorch to within 1.2e-5 in the logits. Other\n", - "encoders and architectures have not been run in Earth Engine as whole models. And a supported\n", - "model is not necessarily an accurate one: that depends on its training.\n", - "\n", - "Encoders checked offline in the same way: ResNet-18, 34 and 50 and VGG-11 (with BatchNorm), in\n", - "`Unet` and `UnetPlusPlus`." + "- **\"Supported\"** means TensorEE can express the model; it is not yet a claim about the numbers.\n", + "- **Evidence so far:** the blocks agree with PyTorch at their validated settings, and one trained\n", + " ResNet-18 U-Net prepared this way matched PyTorch in Earth Engine to 1.2e-5 (64 × 64 input).\n", + " Other encoders have not been run live as whole models.\n", + "- **Checked offline the same way:** ResNet-18, 34, 50 and VGG-11 (BatchNorm) encoders in `Unet`\n", + " and `UnetPlusPlus`.\n", + "- Supported is not accurate: accuracy depends on training." ] }, { @@ -178,8 +160,7 @@ "\n", "### A LinkNet: transposed convolutions\n", "\n", - "LinkNet's decoder upsamples with transposed convolutions (`ConvTranspose2d`). TensorEE has no\n", - "building block for them in a CNN, so it refuses the model and names each one." + "- LinkNet's decoder upsamples with `ConvTranspose2d`, which has no CNN building block." ] }, { @@ -200,10 +181,9 @@ "source": [ "### DeepLabV3+: dilated convolutions, and more\n", "\n", - "DeepLabV3+ keeps its last encoder stage at a higher resolution by spacing out the taps of its\n", - "convolutions (dilation). The CNN path supports convolutions with stride 1 or 2, no dilation,\n", - "and \"same\" zero padding, so this model is refused. The report also names its other missing\n", - "operations: the adaptive average pooling and the bilinear resize of its pyramid pooling module." + "- Its encoder uses dilated convolutions; the CNN path supports stride 1 or 2, no dilation,\n", + " \"same\" zero padding.\n", + "- The report also names its adaptive average pooling and bilinear resize." ] }, { @@ -222,16 +202,14 @@ "id": "f73d6c10", "metadata": {}, "source": [ - "A refusal is a firm answer: TensorEE never replaces an operation it cannot compute exactly with\n", - "an approximation. The message names the first operations that block the model; the\n", - "operations after them are not checked. Your options:\n", - "\n", - "- **Change the architecture**, for example a U-Net decoder (nearest-neighbour upsampling and\n", - " convolutions) in place of transposed convolutions, and retrain.\n", - "- **Ask for the operation.** [Open an issue](https://github.com/noelEOS/tensor-ee/issues)\n", - " with the printed message and your model's architecture. Adding a building block is a\n", - " library change with its own tests and evidence;\n", - " [CONTRIBUTING](https://github.com/noelEOS/tensor-ee/blob/main/CONTRIBUTING.md) describes it." + "**A refusal is firm:** TensorEE never approximates. It names the blocking operations; the ones\n", + "after them are not checked. Options:\n", + "\n", + "- **Change the architecture** (for example a U-Net decoder instead of transposed convolutions)\n", + " and retrain.\n", + "- **Ask for the operation:** [open an issue](https://github.com/noelEOS/tensor-ee/issues) with\n", + " the message and your architecture. A new block comes with its own tests and evidence\n", + " ([CONTRIBUTING](https://github.com/noelEOS/tensor-ee/blob/main/CONTRIBUTING.md))." ] }, { @@ -241,10 +219,9 @@ "source": [ "## 5. Attention models: the converter\n", "\n", - "Transformer-style models (SegFormer, Swin) do not go through `adapt_smp_model`. It refuses them\n", - "and points to the converter, `tensor_ee.converter.convert`, which recognizes the modules of the\n", - "models it was built for: SegFormer-B0 and a UPerNet with a Swin-T encoder, both on\n", - "224 × 224 chips." + "- `adapt_smp_model` refuses transformer models and points to the converter,\n", + " `tensor_ee.converter.convert`.\n", + "- The converter was built for SegFormer-B0 and a Swin-T UPerNet, on 224 × 224 chips." ] }, { @@ -270,12 +247,10 @@ "id": "01ed7f1c", "metadata": {}, "source": [ - "The converter needs the chip's grid. Converting does not compute anything on imagery: it runs\n", - "the model's Python code with placeholder tensors to record its structure. (That also means it\n", - "runs the model's code, so convert only models whose source you trust.)\n", - "\n", - "Its report gives one verdict per building block: **inside** the settings that were validated in\n", - "Earth Engine, or **outside** them." + "- Converting records the model's structure by running its code on placeholder tensors, so\n", + " convert only models whose source you trust.\n", + "- The report gives each block a verdict: **inside** the settings validated in Earth Engine, or\n", + " **outside**." ] }, { @@ -295,12 +270,10 @@ "id": "4fbb3f2c", "metadata": {}, "source": [ - "All 18 blocks of SegFormer-B0 are inside their validated records, and the two \"layout\n", - "operations\" only move pixels onto the model's internal grid and back.\n", - "\n", - "SegFormer-B1 has the same structure at larger widths. It converts, but its blocks are\n", - "**outside** their records: the library can express it, and nothing validates it yet. An\n", - "outside verdict is reported, not refused." + "- SegFormer-B0: all 18 blocks inside; the 2 layout operations only move pixels on and off the\n", + " model's grid.\n", + "- SegFormer-B1 (same structure, wider) converts, but its blocks are **outside**: expressible,\n", + " not yet validated. Outside is reported, not refused." ] }, { @@ -321,9 +294,8 @@ "id": "82600b1d", "metadata": {}, "source": [ - "And a refusal: SegFormer-B0 on a 512 × 512 chip. At that size SegFormer's attention reduces\n", - "its keys and values to a grid larger than one 7 × 7 window, which needs a building block the\n", - "library does not have. The converter says so with a code and the module it stopped at." + "- A refusal: SegFormer-B0 at 512 × 512 reduces its keys and values to more than one 7 × 7\n", + " window, which needs a block the library does not have." ] }, { @@ -347,9 +319,8 @@ "source": [ "## 6. Your model\n", "\n", - "Put your model in `MY_MODEL` and its chip shape in `MY_INPUT_SHAPE`, then run the next two\n", - "cells. Build the architecture exactly as you trained it, and load your weights if you have\n", - "them (random weights give the same answer). For example:\n", + "- Put your model in `MY_MODEL`, built exactly as trained (random weights give the same answer),\n", + " and its chip shape in `MY_INPUT_SHAPE`. For example:\n", "\n", "```python\n", "MY_MODEL = smp.Unet(\"resnet34\", encoder_weights=None, in_channels=4, classes=3)\n", @@ -357,9 +328,7 @@ "MY_INPUT_SHAPE = (1, 4, 256, 256)\n", "```\n", "\n", - "The check runs locally. It tries the CNN path first (for a segmentation-models-pytorch model,\n", - "or a plain PyTorch CNN), and if that does not support the model, the converter for attention\n", - "models." + "- The check tries the CNN path (SMP or plain PyTorch CNN), then the converter." ] }, { @@ -432,20 +401,16 @@ "source": [ "## What next\n", "\n", - "- **Supported, a CNN.** [Notebook 03](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", - " takes a trained segmentation-models-pytorch U-Net from its checkpoint to Earth Engine: it\n", - " plans the tiles, uploads the weights once, predicts one tile, compares it with PyTorch and\n", - " exports a region.\n", - "- **Converted, an attention model.** The\n", - " [converter guide](https://github.com/noelEOS/tensor-ee/blob/main/docs/converter.md) plans the\n", - " model into stages that fit Earth Engine's limits and runs them.\n", - "- **Outside the validated records.** The model can be expressed, but its numbers have no\n", - " evidence yet. Running it and comparing it with PyTorch on your own input is how that evidence\n", - " is made; the maintainer can advise.\n", - "- **Refused.** See the options under step 4.\n", - "\n", - "The [block catalogue](https://github.com/noelEOS/tensor-ee/blob/main/docs/building-blocks.md)\n", - "lists every building block with the settings it was validated at." + "- **Supported CNN:** [notebook 03](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", + " takes it from checkpoint to an exported region.\n", + "- **Converted attention model:** the\n", + " [converter guide](https://github.com/noelEOS/tensor-ee/blob/main/docs/converter.md) plans and\n", + " runs its stages.\n", + "- **Outside the validated records:** run it and compare with PyTorch on your input to make that\n", + " evidence; the maintainer can advise.\n", + "- **Refused:** see step 4. The\n", + " [block catalogue](https://github.com/noelEOS/tensor-ee/blob/main/docs/building-blocks.md)\n", + " lists every block and its validated settings." ] } ], diff --git a/notebooks/03_bring_your_own_cnn.ipynb b/notebooks/03_bring_your_own_cnn.ipynb index 01fbc39..851799d 100644 --- a/notebooks/03_bring_your_own_cnn.ipynb +++ b/notebooks/03_bring_your_own_cnn.ipynb @@ -9,47 +9,32 @@ "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", "\n", - "This notebook takes a U-Net you trained with\n", - "[segmentation-models-pytorch](https://github.com/qubvel-org/segmentation_models.pytorch)\n", - "(SMP) and runs it inside Earth Engine on your own imagery. The steps:\n", - "\n", - "1. load your checkpoint and check that TensorEE supports the model;\n", - "2. work out how much extra border (the halo) each tile needs, and plan the tiles;\n", - "3. upload the model's weights to Earth Engine once, as table assets;\n", - "4. predict one tile and compare it with PyTorch on the same pixels;\n", - "5. export a whole region.\n", - "\n", - "**Offline first.** With the settings as they are, the notebook uses a small demonstration\n", - "U-Net with random weights and runs every step offline: it builds the Earth Engine\n", - "computations but sends nothing. Then you set your checkpoint, your project and your input,\n", - "and turn on the billed steps one at a time. Each billed step says beforehand what it creates\n", - "and what is known of its cost, and runs only when both `RUN_LIVE` and its own switch are\n", - "`True`.\n", - "\n", - "**What has been checked.** Each building block this U-Net uses agrees with PyTorch at the\n", - "settings it was validated at. One trained U-Net with a ResNet-18 encoder, prepared the way\n", - "this notebook prepares yours, ran in Earth Engine on a 64 × 64 test input and matched PyTorch\n", - "within 1.2e-5 in the logits. The tiled recipe in this notebook has **not** been run in Earth\n", - "Engine as a whole: step 9's comparison is how you check it for your model. Agreement with\n", - "PyTorch says nothing about accuracy, which you measure on your own labelled data." - ] - }, - { - "cell_type": "markdown", - "id": "b29a1f17", - "metadata": {}, - "source": [ - "## 1. Install\n", - "\n", - "The `segwater` extra brings segmentation-models-pytorch 0.5.0, the version the library was\n", - "checked with. PyTorch is already installed in Colab." + "Takes a U-Net trained with\n", + "[segmentation-models-pytorch](https://github.com/qubvel-org/segmentation_models.pytorch) (SMP)\n", + "and runs it inside Earth Engine on your imagery:\n", + "\n", + "1. check the model and find its **halo** (the border each tile needs);\n", + "2. plan the tiles;\n", + "3. upload the weights once, as Earth Engine table assets;\n", + "4. predict one test tile and compare it with PyTorch;\n", + "5. export a region.\n", + "\n", + "- **Offline first.** As shipped, it uses a small demo U-Net with random weights and only builds\n", + " the Earth Engine computations; nothing is sent.\n", + "- **Then live, step by step.** Set your checkpoint, project and input, set `RUN_LIVE = True`,\n", + " and turn on each Earth Engine step with its own switch. Each step uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate.\n", + "- **What is checked.** The building blocks agree with PyTorch at their validated settings, and\n", + " one trained ResNet-18 U-Net, prepared this way, matched PyTorch to 1.2e-5 in Earth Engine on a\n", + " 64 × 64 input. This tiled recipe has not been run live as a whole: step 9 checks it for your\n", + " model. Agreement with PyTorch is not accuracy." ] }, { "cell_type": "code", - "id": "b0fec902", + "id": "b29a1f17", "metadata": {}, "source": [ + "# @title 1. Install\n", "%pip install --quiet \"tensor-ee[segwater]\"" ], "execution_count": null, @@ -57,60 +42,49 @@ }, { "cell_type": "markdown", - "id": "2f178b02", + "id": "b0fec902", "metadata": {}, "source": [ "## 2. Settings\n", "\n", - "Edit this cell, and only this cell.\n", - "\n", - "**Your model.** `CHECKPOINT` is the path of your trained weights, saved with\n", - "`torch.save(model.state_dict(), path)`. In Colab, upload it with the **Files** panel on the left.\n", - "Leave it `None` to use the small demonstration model, which runs offline only. `ENCODER`,\n", - "`BANDS` and `CLASSES` must match how you trained the model.\n", - "\n", - "**Your input.** `INPUT_IMAGE` is an Earth Engine image with the bands `BANDS`, already scaled\n", - "the way your training data was (TensorEE does not normalize for you). `CRS`, `REGION_X`,\n", - "`REGION_Y` and `PIXEL` place the region: the upper-left corner of its first tile, in metres of\n", - "`CRS`. The image must also cover a border of one halo around the region (step 5 prints how\n", - "wide).\n", - "\n", - "**The billed steps.** Each runs only when `RUN_LIVE` and its own switch are both `True`.\n", - "\n", - "After changing `RUN_LIVE`, restart the runtime (in Colab: **Runtime → Restart session**) and\n", - "run the notebook again from the top: the offline mode stays installed until the runtime\n", - "restarts." + "- **Model:** `CHECKPOINT` is your `state_dict` file (in Colab, upload it in the **Files** panel).\n", + " `None` = the offline demo. `ENCODER`, `BANDS` and `CLASSES` must match your training.\n", + "- **Input:** `INPUT_IMAGE` is an Earth Engine image with bands `BANDS`, already scaled as in\n", + " training. `REGION_X`, `REGION_Y` are the region's upper-left corner in `CRS` metres; the image\n", + " must also cover one halo around it (step 5 prints the width).\n", + "- **Earth Engine steps** run only when `RUN_LIVE` and their own switch are both `True`. You can\n", + " change `RUN_LIVE` at any time and run from the top again." ] }, { "cell_type": "code", - "id": "6193d146", + "id": "2f178b02", "metadata": {}, "source": [ "PROJECT = \"YOUR_EE_PROJECT\" # your Cloud project registered for Earth Engine\n", - "RUN_LIVE = False # True: sign in; the switches below then choose the billed steps\n", + "RUN_LIVE = False # True: sign in; the switches below choose the Earth Engine steps\n", "\n", "# Your model\n", - "CHECKPOINT = None # for example \"my_unet.pth\"; None = demonstration model, offline only\n", + "CHECKPOINT = None # for example \"my_unet.pth\"; None = demo model, offline only\n", "ENCODER = \"resnet34\"\n", - "BANDS = (\"VV\", \"VH\") # your input bands, in training order\n", + "BANDS = (\"VV\", \"VH\") # input bands, in training order\n", "CLASSES = 2\n", "\n", - "# Your input and the region\n", + "# Your input and region\n", "INPUT_IMAGE = \"projects/YOUR_EE_PROJECT/assets/YOUR_PREPARED_INPUT\"\n", "CRS = \"EPSG:32651\"\n", - "REGION_X, REGION_Y = 286030, 2640860 # upper-left corner of the region, in CRS metres\n", + "REGION_X, REGION_Y = 286030, 2640860 # upper-left corner of the region, CRS metres\n", "PIXEL = 10 # metres per pixel\n", - "TILE = 1024 # tile side for the region, in pixels\n", - "REGION_COLUMNS, REGION_ROWS = 2, 1 # the region is this many tiles\n", - "TEST_TILE = 256 # side of the single test tile of step 8\n", + "TILE = 1024 # region tile side, pixels\n", + "REGION_COLUMNS, REGION_ROWS = 2, 1 # region size, in tiles\n", + "TEST_TILE = 256 # side of the test tile (step 8)\n", "\n", "# What the live steps create\n", - "WEIGHTS_MANIFEST = \"my_unet_weights.json\" # a local file recording the uploaded weights\n", - "TEST_ASSET = \"my_unet_test_tile\" # new image assets in your project\n", + "WEIGHTS_MANIFEST = \"my_unet_weights.json\" # local record of the uploaded weights: keep it\n", + "TEST_ASSET = \"my_unet_test_tile\"\n", "REGION_ASSET = \"my_unet_region\"\n", "\n", - "# The billed steps\n", + "# Earth Engine steps\n", "UPLOAD_WEIGHTS = False # step 7\n", "PREDICT_TEST_TILE = False # step 8\n", "COMPARE_TEST_TILE = False # step 9\n", @@ -121,96 +95,80 @@ }, { "cell_type": "code", - "id": "2cfc7c5f", + "id": "6193d146", "metadata": {}, "source": [ - "PLACEHOLDER = \"YOUR_EE_PROJECT\"\n", + "class Stop(Exception):\n", + " \"\"\"Stops the notebook with a one-line message instead of a traceback.\"\"\"\n", + "\n", + " def _render_traceback_(self):\n", + " return [f\"Stopped: {self}\"]\n", "\n", "\n", "def live_step(what, switch=True):\n", - " \"\"\"True when a billed step should run; refuses a placeholder project or a missing checkpoint.\"\"\"\n", + " \"\"\"True when an Earth Engine step should run.\"\"\"\n", " if not (RUN_LIVE and switch):\n", - " reason = \"RUN_LIVE is False\" if not RUN_LIVE else \"its switch is False\"\n", - " print(f\"Skipped: {what}. {reason}, so nothing was sent to Earth Engine.\")\n", + " print(f\"Skipped: {what} ({'RUN_LIVE' if not RUN_LIVE else 'its switch'} is False).\")\n", " return False\n", - " if not PROJECT or PROJECT == PLACEHOLDER:\n", - " raise ValueError(\"Set PROJECT to your own Cloud project ID in the settings cell.\")\n", + " if not PROJECT or PROJECT == \"YOUR_EE_PROJECT\":\n", + " raise Stop(\"set PROJECT to your Cloud project ID in the settings cell.\")\n", " if CHECKPOINT is None:\n", - " raise ValueError(\"Live steps need your trained checkpoint: set CHECKPOINT. \"\n", - " \"The demonstration model has random weights.\")\n", + " raise Stop(\"live steps need your trained checkpoint: set CHECKPOINT.\")\n", " return True" ], "execution_count": null, "outputs": [] }, { - "cell_type": "markdown", - "id": "695f4df2", + "cell_type": "code", + "id": "2cfc7c5f", "metadata": {}, "source": [ - "## 3. Sign in, or work offline\n", + "# @title 3. Connect: offline signatures, or your project\n", + "# @markdown Offline, Earth Engine's own test helper lets the notebook build computations without\n", + "# @markdown credentials. Switching to live undoes it, so no restart is needed.\n", + "import ee\n", + "from ee.apitestcase import ApiTestCase\n", "\n", - "With `RUN_LIVE = True` the first cell signs you in; the first time, `ee.Authenticate()` shows a\n", - "link or a prompt. Signing in costs nothing by itself. With `RUN_LIVE = False` the second cell\n", - "installs the Earth Engine client's offline function signatures (`ee.apitestcase.ApiTestCase`,\n", - "a test helper that ships with the client), so the notebook can build Earth Engine computations\n", - "without credentials and without sending anything." - ] + "if \"offline\" in globals(): # an earlier offline run in this session: undo it\n", + " offline.tearDown()\n", + " ee.data._install_cloud_api_resource = _install_cloud_api_resource\n", + " ee.Reset()\n", + " del offline\n", + "if not RUN_LIVE:\n", + " _install_cloud_api_resource = ee.data._install_cloud_api_resource\n", + " offline = ApiTestCase()\n", + " offline.setUp()\n", + " print(\"Offline: computations are built, nothing is sent.\")" + ], + "execution_count": null, + "outputs": [] }, { "cell_type": "code", - "id": "84bcd7d7", + "id": "695f4df2", "metadata": { "tags": [ "live" ] }, "source": [ - "import ee\n", - "\n", - "if live_step(\"connect to Earth Engine\"):\n", + "if live_step(\"sign in\"):\n", " ee.Authenticate()\n", " ee.Initialize(project=PROJECT)\n", - " print(\"Signed in. Requests are billed to\", PROJECT)" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "d4b2bff3", - "metadata": {}, - "source": [ - "import ee\n", - "\n", - "if not RUN_LIVE:\n", - " from ee.apitestcase import ApiTestCase\n", - "\n", - " offline = ApiTestCase()\n", - " offline.setUp() # offline signatures: no credentials, no network\n", - " print(\"Offline: the notebook builds Earth Engine computations but sends none.\")" + " print(\"Signed in to\", PROJECT)" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "aa4d0e83", - "metadata": {}, - "source": [ - "## 4. Load the model\n", - "\n", - "With `CHECKPOINT = None` this builds a shallow U-Net (three encoder stages, small decoder)\n", - "with random weights, so that the offline steps stay quick. With your checkpoint, it builds the\n", - "standard U-Net and loads your weights strictly: a missing or unexpected weight name stops it,\n", - "which usually means the settings do not match your training." - ] - }, { "cell_type": "code", - "id": "42d821f2", + "id": "84bcd7d7", "metadata": {}, "source": [ + "# @title 4. Load the model\n", + "# @markdown `None`: a shallow demo U-Net with random weights. A checkpoint is loaded strictly:\n", + "# @markdown a missing or extra weight name means the settings do not match your training.\n", "import torch\n", "import segmentation_models_pytorch as smp\n", "\n", @@ -227,111 +185,80 @@ "if CHECKPOINT is None:\n", " torch.manual_seed(20261020)\n", " model = make_model(encoder_depth=3, decoder_channels=(16, 8, 4))\n", - " tile, test_tile = 64, 64 # small tiles for the demonstration\n", - " print(\"Demonstration model: a shallow U-Net with random weights.\")\n", + " tile, test_tile = 64, 64 # small tiles for the demo\n", + " print(\"Demo model: a shallow U-Net with random weights.\")\n", "else:\n", " model = make_model()\n", " model.load_state_dict(torch.load(CHECKPOINT, map_location=\"cpu\", weights_only=True), strict=True)\n", " tile, test_tile = TILE, TEST_TILE\n", - " print(\"Loaded\", CHECKPOINT)\n", - "print(f\"{sum(p.numel() for p in model.parameters()):,} parameters\")" + " print(\"Loaded\", CHECKPOINT)" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "73b9c3ce", - "metadata": {}, - "source": [ - "## 5. Check support, and find the halo\n", - "\n", - "`adapt_smp_model` prepares the model for TensorEE at a fixed input size, and checks that the\n", - "prepared model gives exactly the same numbers as yours on a test input. `check_compatibility`\n", - "then checks every operation. (Notebook 02 explains these reports in more detail.)\n", - "\n", - "`analyze_receptive_field` works out how far each output pixel can see into the input. That\n", - "reach is the **halo**: to compute a tile exactly as the model would compute it on a larger\n", - "image, Earth Engine must read the tile plus a halo on every side, and the halo is thrown away\n", - "afterwards. For a standard U-Net with a ResNet-34 encoder the halo is 544 pixels, so a\n", - "1,024-pixel tile reads 2,112 × 2,112 pixels: about 4.25 input pixels per output pixel. Larger\n", - "tiles waste less." - ] - }, { "cell_type": "code", - "id": "3c1fd62e", + "id": "d4b2bff3", "metadata": {}, "source": [ + "# @title 5. Check support and find the halo\n", + "# @markdown The halo is how far each output pixel sees into the input. Every tile reads the\n", + "# @markdown tile plus a halo on each side, then drops it: for a ResNet-34 U-Net, 544 pixels,\n", + "# @markdown so a 1,024-pixel tile reads 2,112 × 2,112. Larger tiles waste less.\n", "from tensor_ee import check_compatibility\n", "from tensor_ee.receptive_field import analyze_receptive_field\n", "from tensor_ee.smp import adapt_smp_model\n", "\n", "\n", "def prepare(height, width):\n", - " \"\"\"The model adapted to one input size, refused unless every operation is supported.\"\"\"\n", + " \"\"\"The model adapted to one input size; stops unless every operation is supported.\"\"\"\n", " shape = (1, len(BANDS), height, width)\n", " adapted = adapt_smp_model(model, input_shape=shape)\n", " report = check_compatibility(adapted, shape)\n", " if report.status != \"supported\":\n", - " raise ValueError(f\"TensorEE does not support this model: {report.to_dict()['blockers'][:3]}\")\n", + " raise Stop(f\"TensorEE does not support this model: {report.to_dict()['blockers'][:3]} \"\n", + " \"(notebook 02 explains reports).\")\n", " return adapted\n", "\n", "\n", "probe = prepare(64, 64)\n", "support = analyze_receptive_field(probe)\n", "halo, alignment = support.safe_halo, support.alignment\n", - "print(\"Supported. Prepared model equals the original on a test input:\",\n", + "print(\"Supported; the prepared model equals yours on a test input:\",\n", " probe.tensor_ee_smp_metadata[\"torch_equal\"])\n", - "print(f\"Halo: {halo} pixels on each side ({halo * PIXEL / 1000:g} km). \"\n", - " f\"Tile sizes and positions must be multiples of {alignment} pixels.\")" + "print(f\"Halo {halo} px ({halo * PIXEL / 1000:g} km) per side; sizes in multiples of {alignment} px.\")" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "41ea3bf8", - "metadata": {}, - "source": [ - "## 6. Plan the tiles\n", - "\n", - "The plan puts the region's tiles side by side and gives each one its halo. Every tile then\n", - "reads an input of the same size, which the model needs. The region's input must cover the\n", - "region plus the halo on every side.\n", - "\n", - "There are two plans: one test tile, for step 8, and the region, for step 10." - ] - }, { "cell_type": "code", - "id": "1a50cc0b", + "id": "aa4d0e83", "metadata": {}, "source": [ + "# @title 6. Plan the tiles\n", + "# @markdown One test tile (step 8) and the region (step 10), each with a full halo of real input.\n", "from tensor_ee import Grid\n", "from tensor_ee.converter import PixelWindow, plan_tiles\n", "\n", "\n", "def plan_region(side, columns, rows):\n", - " \"\"\"Tiles of side x side pixels, with a full halo of real input around the whole region.\"\"\"\n", " for name, value in ((\"tile side\", side), (\"halo\", halo)):\n", " if value % alignment:\n", - " raise ValueError(f\"The {name} ({value}) must be a multiple of {alignment} pixels.\")\n", - " width, height = columns * side + 2 * halo, rows * side + 2 * halo\n", + " raise Stop(f\"the {name} ({value}) must be a multiple of {alignment} pixels.\")\n", " padded = Grid(CRS, (PIXEL, 0, REGION_X - halo * PIXEL, 0, -PIXEL, REGION_Y + halo * PIXEL),\n", - " width, height)\n", + " columns * side + 2 * halo, rows * side + 2 * halo)\n", " region = PixelWindow(halo, halo, columns * side, rows * side)\n", " return plan_tiles(padded, side, halo, alignment=alignment, roi=region)\n", "\n", "\n", "def describe(plan, name):\n", " context = plan.tiles[0].context\n", - " read = sum(t.context.pixels for t in plan.tiles)\n", - " print(f\"{name}: {len(plan.tiles)} tile(s) of {plan.core_size[0]} x {plan.core_size[1]} pixels; \"\n", - " f\"each reads {context.width} x {context.height}; \"\n", - " f\"{read / plan.roi.pixels:.2f} input pixels per output pixel.\")\n", - " print(f\" Input needed: {plan.global_grid.width} x {plan.global_grid.height} pixels \"\n", - " f\"from corner {plan.global_grid.transform[2]}, {plan.global_grid.transform[5]} in {CRS}.\")\n", + " ratio = sum(t.context.pixels for t in plan.tiles) / plan.roi.pixels\n", + " print(f\"{name}: {len(plan.tiles)} tile(s) of {plan.core_size[0]} px, each reading \"\n", + " f\"{context.width} x {context.height} ({ratio:.2f} input px per output px); input \"\n", + " f\"{plan.global_grid.width} x {plan.global_grid.height} px from \"\n", + " f\"({plan.global_grid.transform[2]}, {plan.global_grid.transform[5]}).\")\n", "\n", "\n", "test_plan = plan_region(test_tile, 1, 1)\n", @@ -342,34 +269,16 @@ "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "47fb0596", - "metadata": {}, - "source": [ - "## 7. Upload the weights once (billed to your project)\n", - "\n", - "TensorEE folds each BatchNorm layer into the convolution before it (the same arithmetic, fewer\n", - "steps) and stores the resulting numbers in Earth Engine table assets. Every prediction then\n", - "reads them from there, so they are uploaded once and reused.\n", - "\n", - "**What this creates.** Table assets in your project, under a new name\n", - "`projects/PROJECT/assets/tensor_ee_weights_`, each holding at most 8 MB of numbers as\n", - "text; the cell prints an estimate of how many. One batch table-export task per table, at most\n", - "two at a time; the cell waits for all of them. Locally, the file `WEIGHTS_MANIFEST`, which\n", - "records the tables and a fingerprint of every convolution. Keep that file: the later steps\n", - "need it, and in Colab it disappears when the runtime ends, so download it from the **Files**\n", - "panel.\n", - "\n", - "**What it costs.** Not recorded. The trained ResNet-18 U-Net of notebook 01 has 14.3 million\n", - "numbers and took 29 tables." - ] - }, { "cell_type": "code", - "id": "d6897497", + "id": "42d821f2", "metadata": {}, "source": [ + "# @title 7. Upload the weights once\n", + "# @markdown - BatchNorm is folded into the convolutions, and the numbers go into table assets\n", + "# @markdown `projects/PROJECT/assets/tensor_ee_weights_`, one export task each (two at a time).\n", + "# @markdown - `WEIGHTS_MANIFEST` records them locally: keep it (in Colab, download it).\n", + "# @markdown - Cost not recorded; it uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate. The trained ResNet-18 U-Net (14.3 M numbers) took 29 tables.\n", "import math\n", "\n", "from tensor_ee import effective_convolutions\n", @@ -384,7 +293,7 @@ }, { "cell_type": "code", - "id": "1aaa8523", + "id": "73b9c3ce", "metadata": { "tags": [ "live" @@ -401,8 +310,7 @@ "if live_step(\"upload the weights\", UPLOAD_WEIGHTS):\n", " manifest_path = Path(WEIGHTS_MANIFEST)\n", " if manifest_path.exists():\n", - " raise FileExistsError(f\"{manifest_path} already exists: the weights are uploaded. \"\n", - " \"Use a new WEIGHTS_MANIFEST name only for new weights.\")\n", + " raise Stop(f\"{manifest_path} exists, so the weights are already uploaded.\")\n", " namespace = f\"projects/{PROJECT}/assets/tensor_ee_weights_{uuid.uuid4().hex[:12]}\"\n", "\n", " def progress(manifest):\n", @@ -410,32 +318,20 @@ " print(json.dumps({\"status\": manifest[\"status\"], \"tables\": dict(states)}), flush=True)\n", "\n", " export_weights(parameters, namespace, manifest_path, on_update=progress)\n", - " print(\"Weights uploaded; manifest written to\", manifest_path)" + " print(\"Uploaded; manifest written to\", manifest_path)" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "09c5f873", - "metadata": {}, - "source": [ - "## 8. Predict the test tile (billed to your project)\n", - "\n", - "The next cell builds the prediction: the model, compiled to Earth Engine operations for the test tile's input size, with\n", - "its weights read from the uploaded tables. Offline, the tables are placeholders that do not\n", - "exist, which is enough to build and measure the computation.\n", - "\n", - "The result has these bands: `score_0`, `score_1`, … (the raw logits), `p_0`, `p_1`, …\n", - "(probabilities), `class_index` (the most likely class) and, for two classes, `binary_class`\n", - "(1 where `p_1` ≥ 0.5). A prediction is made only where every input band is present." - ] - }, { "cell_type": "code", - "id": "0c458253", + "id": "3c1fd62e", "metadata": {}, "source": [ + "# @title 8. Predict the test tile\n", + "# @markdown Builds the computation (offline, with stand-in weight tables). Output bands: `score_k`\n", + "# @markdown (logits), `p_k` (probabilities), `class_index` and, for two classes, `binary_class`\n", + "# @markdown (`p_1` ≥ 0.5). Pixels missing any input band get no prediction.\n", "import json\n", "from pathlib import Path\n", "\n", @@ -448,11 +344,10 @@ "\n", "\n", "def weights_manifest():\n", - " \"\"\"The uploaded weights, or offline a stand-in that names tables which do not exist.\"\"\"\n", " if RUN_LIVE and Path(WEIGHTS_MANIFEST).exists():\n", " return json.loads(Path(WEIGHTS_MANIFEST).read_text())\n", " if RUN_LIVE:\n", - " raise FileNotFoundError(f\"{WEIGHTS_MANIFEST} not found: upload the weights first (step 7).\")\n", + " raise Stop(f\"{WEIGHTS_MANIFEST} not found: upload the weights first (step 7).\")\n", " return {\"format\": \"native_ee_featurecollection_weights_v2\", \"status\": \"ready\",\n", " \"layers\": {p.key: {\"shape\": list(p.shape), \"dtype\": p.dtype, \"bias\": p.bias is not None,\n", " \"source\": p.source, \"sha256\": p.sha256,\n", @@ -467,11 +362,10 @@ "\n", "\n", "def build_prediction(plan):\n", - " \"\"\"The model's Earth Engine computation over a plan's tiles, with its output bands.\"\"\"\n", " context = plan.tiles[0].context\n", " adapted = prepare(context.height, context.width)\n", " if analyze_receptive_field(adapted).safe_halo != halo:\n", - " raise ValueError(\"The halo changed with the input size; revise the plan.\")\n", + " raise Stop(\"the halo changed with the input size; revise the plan.\")\n", " provider = MatchedWeightProvider(adapted, weights_manifest())\n", " source = input_image().select(list(BANDS)).toFloat().reproject(plan.global_grid.projection())\n", " valid = Raster(source.mask().reduce(ee.Reducer.min()).rename(\"valid\"), (\"valid\",), plan.global_grid)\n", @@ -501,86 +395,58 @@ "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "757fef81", - "metadata": {}, - "source": [ - "**What the export creates.** One batch export task and, when it completes, one new image\n", - "asset `projects/PROJECT/assets/TEST_ASSET` over the test tile (not its halo).\n", - "\n", - "**What it costs.** Not measured for this recipe. For scale: the trained water U-Net's\n", - "224 × 224 chip took 64.9 batch EECU-seconds, and this test tile reads a larger input\n", - "(step 6). If Earth Engine reports that it ran out of memory or time, reduce `TEST_TILE` and\n", - "try again; do not just resubmit." - ] - }, { "cell_type": "code", - "id": "df5fae6a", + "id": "41ea3bf8", "metadata": { "tags": [ "live" ] }, "source": [ + "# @title Export the test tile\n", + "# @markdown One batch export, one new image asset (the tile, not its halo). Cost not measured;\n", + "# @markdown it uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate. For scale, notebook 01's 224 × 224 chip used about 65 EECU-seconds.\n", + "# @markdown If it runs out of memory or time, lower `TEST_TILE` rather than resubmitting.\n", "import time\n", "\n", "\n", "def wait(task, minutes=60):\n", - " \"\"\"Check a task every 30 seconds until it ends.\"\"\"\n", " deadline = time.time() + minutes * 60\n", " while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", " if time.time() > deadline:\n", - " raise TimeoutError(\"Still running; run this cell again later, or watch the Tasks tab.\")\n", + " raise Stop(\"still running; follow it in the Tasks tab, or run this cell again later.\")\n", " time.sleep(30)\n", - " print(\"Task\", task.id, state)\n", " if state != \"COMPLETED\":\n", - " raise RuntimeError(task.status().get(\"error_message\", \"the export did not complete\"))\n", + " raise Stop(f\"the export ended {state}: {task.status().get('error_message', '')}\")\n", + " print(\"Task\", task.id, \"completed\")\n", "\n", "\n", - "if live_step(\"export the test tile (1 batch task, 1 new asset)\", PREDICT_TEST_TILE):\n", + "if live_step(\"export the test tile\", PREDICT_TEST_TILE):\n", " test_task = export_task(test_prediction, test_plan, TEST_ASSET, \"tensor_ee_test_tile\")\n", " test_task.start()\n", - " print(\"Started\", test_task.id)\n", " wait(test_task)" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "3d569881", - "metadata": {}, - "source": [ - "## 9. Compare the test tile with PyTorch (billed to your project)\n", - "\n", - "This is the check that matters for your model: Earth Engine's logits against your model's in\n", - "PyTorch, on exactly the same input pixels. The cell downloads the test tile's input (tile plus\n", - "halo) and the exported logits, runs your model on the input locally, and reports the largest\n", - "absolute difference. TensorEE's target is 1e-4.\n", - "\n", - "**What this does.** Two interactive requests (`ee.data.computePixels`): the input and the\n", - "exported result. It creates nothing. The input must be complete over the tile and its halo.\n", - "\n", - "**What it costs.** Small interactive requests, not measured. The cell refuses an input over\n", - "40 MB." - ] - }, { "cell_type": "code", - "id": "b3f6e6aa", + "id": "1a50cc0b", "metadata": { "tags": [ "live" ] }, "source": [ + "# @title 9. Compare the test tile with PyTorch\n", + "# @markdown Downloads the tile's input (with halo) and the exported logits (2 small requests,\n", + "# @markdown no assets), runs your model locally, and reports the largest absolute difference\n", + "# @markdown (target 1e-4). The input must be complete; it is refused above 40 MB.\n", "import numpy as np\n", "\n", "\n", "def download(image, grid, bands):\n", - " \"\"\"The exact float32 pixels of an image's bands on a grid: one interactive request.\"\"\"\n", " names = (\"scaleX\", \"shearX\", \"translateX\", \"shearY\", \"scaleY\", \"translateY\")\n", " values = ee.data.computePixels({\n", " \"expression\": image.select(list(bands)), \"fileFormat\": \"NUMPY_NDARRAY\",\n", @@ -589,64 +455,48 @@ " return np.stack([values[band] for band in bands]).astype(np.float32)\n", "\n", "\n", - "if live_step(\"compare the test tile with PyTorch (2 requests)\", COMPARE_TEST_TILE):\n", + "if live_step(\"compare the test tile with PyTorch\", COMPARE_TEST_TILE):\n", " padded = test_plan.global_grid\n", " size = padded.width * padded.height * (len(BANDS) + 1) * 4\n", " if size > 40_000_000:\n", - " raise ValueError(f\"The input is {size / 1e6:.0f} MB; reduce TEST_TILE.\")\n", + " raise Stop(f\"the input is {size / 1e6:.0f} MB; lower TEST_TILE.\")\n", " source = ee.Image(INPUT_IMAGE).select(list(BANDS)).toFloat().reproject(padded.projection())\n", " valid = source.mask().reduce(ee.Reducer.min()).rename(\"valid\")\n", " pixels = download(source.unmask(0).addBands(valid), padded, (*BANDS, \"valid\"))\n", " if pixels[-1].min() < 1:\n", - " raise ValueError(\"The input has missing pixels in the tile or its halo; the comparison \"\n", - " \"needs complete input. Move the region or fill the gaps.\")\n", + " raise Stop(\"the input has gaps in the tile or its halo; move the region or fill them.\")\n", " exported = download(ee.Image(f\"projects/{PROJECT}/assets/{TEST_ASSET}\"), test_plan.roi_grid, SCORES)\n", "\n", " with torch.no_grad():\n", " reference = model(torch.from_numpy(pixels[None, :-1].copy()))[0].numpy()\n", " core = reference[:, halo:halo + test_tile, halo:halo + test_tile]\n", " error = np.abs(exported - core)\n", - " disagreements = int((exported.argmax(0) != core.argmax(0)).sum())\n", - " print(f\"Largest absolute logit difference: {error.max():.3e} (target 1e-4)\")\n", - " print(f\"Mean absolute difference: {error.mean():.3e}; \"\n", - " f\"class disagreements: {disagreements} of {core[0].size} pixels\")\n", + " print(f\"Largest absolute logit difference: {error.max():.3e} (target 1e-4); \"\n", + " f\"mean {error.mean():.3e}; class disagreements \"\n", + " f\"{int((exported.argmax(0) != core.argmax(0)).sum())} of {core[0].size} pixels.\")\n", " print(\"Within the target.\" if error.max() <= 1e-4 else\n", " \"Above the target: please report it, with these numbers, as an issue.\")" ], "execution_count": null, "outputs": [] }, - { - "cell_type": "markdown", - "id": "0713bc1b", - "metadata": {}, - "source": [ - "## 10. Export the region (billed to your project)\n", - "\n", - "**What this creates.** One batch export task and, when it completes, one new image asset\n", - "`projects/PROJECT/assets/REGION_ASSET` over the region (not its halo), with the bands of\n", - "step 8, in float32.\n", - "\n", - "**What it costs.** Not measured. Earth Engine computes every tile with its halo (step 6\n", - "prints how many input pixels it reads per output pixel), and capacity at these tile sizes has\n", - "not been measured: an export can fail for lack of memory. Run the test tile first." - ] - }, { "cell_type": "code", - "id": "15be8dc7", + "id": "47fb0596", "metadata": { "tags": [ "live" ] }, "source": [ - "if live_step(\"export the region (1 batch task, 1 new asset)\", EXPORT_REGION):\n", + "# @title 10. Export the region\n", + "# @markdown One batch export, one new image asset over the region (not its halo). Cost and\n", + "# @markdown capacity not measured at these tile sizes; it uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate. Run the test tile first.\n", + "if live_step(\"export the region\", EXPORT_REGION):\n", " with deep_recursion():\n", - " region_prediction = build_prediction(region_plan)\n", - " region_task = export_task(region_prediction, region_plan, REGION_ASSET, \"tensor_ee_region\")\n", + " region_task = export_task(build_prediction(region_plan), region_plan, REGION_ASSET,\n", + " \"tensor_ee_region\")\n", " region_task.start()\n", - " print(\"Started\", region_task.id, \"- follow it in the Tasks tab, or wait here.\")\n", " wait(region_task, minutes=240)" ], "execution_count": null, @@ -654,24 +504,17 @@ }, { "cell_type": "markdown", - "id": "4cd0dc7c", + "id": "d6897497", "metadata": {}, "source": [ - "## What this shows, and what it does not\n", - "\n", - "- **Offline**, the notebook shows that TensorEE supports the model and can build its Earth\n", - " Engine computation for your tiles.\n", - "- **Step 9** is live agreement for your model, on one tile and one input. It is the evidence\n", - " to keep, with its numbers.\n", - "- **The region** reuses the same computation for every tile; the halo is what makes tile\n", - " borders match a single run of the model, as far as the model's reach goes.\n", - "- **Accuracy** is a separate question: evaluate the exported predictions against your own\n", - " labelled data.\n", - "\n", - "The [guide this notebook follows](https://github.com/noelEOS/tensor-ee/blob/main/docs/own-model.md)\n", - "gives the details: other encoders, irregular regions, and how validity masks are handled.\n", - "**Cleaning up**: the weight tables and image assets stay in your project until you delete them\n", - "in the Code Editor's **Assets** tab; keep the weight tables as long as you predict with them." + "## Good to know\n", + "\n", + "- **Evidence:** step 9 is live agreement for your model on one tile; keep its numbers.\n", + " Accuracy is separate: evaluate the exports against your own labelled data.\n", + "- **Details:** [the guide this notebook follows](https://github.com/noelEOS/tensor-ee/blob/main/docs/own-model.md)\n", + " covers other encoders, irregular regions and validity masks.\n", + "- **Clean up:** weight tables and images stay in your project until you delete them (Code\n", + " Editor, **Assets** tab); keep the weight tables while you predict with them." ] } ], diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py index 08ddb33..7e99427 100644 --- a/tests/test_notebooks.py +++ b/tests/test_notebooks.py @@ -1,10 +1,18 @@ """The reader notebooks in ``notebooks/``: clean, free of private values, and runnable offline. Read with ``json`` (no Jupyter needed). Each notebook's code cells run in order in a fresh -namespace with ``RUN_LIVE = False``, skipping the cells tagged ``live`` and the ``%`` and ``!`` -lines (installs). Every Earth Engine cloud entry point is refused while they run, as in -``demos.code_editor.build.offline_earth_engine``. Notebooks that need segmentation-models-pytorch -or timm skip by name without them. +namespace, skipping the cells tagged ``live`` and the ``%`` and ``!`` lines (installs), with +every Earth Engine data call and ``ee.Authenticate`` refused, as in +``demos.code_editor.build.offline_earth_engine``. + +- The quickstart (01) has no offline mode of its own: the test installs the client's offline + signatures, and after its settings cell supplies a project name and the place and scene that + its live cells would find (the preset Sun Moon Lake chip). +- 03 runs with ``RUN_LIVE = False`` and its own offline signatures; switching it to live must + undo them without a restart. +- 02 is offline throughout. + +Notebooks that need segmentation-models-pytorch or timm skip by name without them. """ from __future__ import annotations @@ -34,6 +42,8 @@ HOME = re.compile(r"/Users/|/home/|[A-Za-z]:\\\\Users|~/") ASSET = re.compile(r"projects/([^/\s\"'`]+)/assets/[^\s\"'`)]*") PROJECT_VALUE = re.compile(r"""project\s*=\s*["']([^"']*)["']""", re.I) +# The quickstart's preset: the Code Editor demo's verified Sun Moon Lake chip. +PRESET = json.loads((ROOT / "demos/code_editor/presets/segwater_sun_moon_lake_1x1.json").read_text()) def _cells(name): @@ -44,6 +54,14 @@ def _source(cell): return "".join(cell["source"]) +def _code(name): + return [cell for cell in _cells(name) if cell["cell_type"] == "code"] + + +def _live(cell): + return "live" in cell["metadata"].get("tags", []) + + def _public_assets(): assets = set() for path in MANIFESTS.glob("*.json"): @@ -81,19 +99,33 @@ def test_each_opens_with_a_colab_badge_for_its_file_on_main(name): @pytest.mark.parametrize("name", sorted(EXPECTED)) -def test_the_settings_start_offline_and_only_live_cells_reach_earth_engine(name): - cells = [cell for cell in _cells(name) if cell["cell_type"] == "code"] +def test_only_live_cells_reach_earth_engine(name): + for cell in _code(name): + if not _live(cell): + assert not LIVE_CALLS.search(_source(cell)), _source(cell)[:200] + + +def test_the_quickstart_has_one_setting_and_stops_cleanly_without_it(): + sources = [_source(cell) for cell in _code("01_quickstart.ipynb")] + text = "\n".join(sources) + assert "RUN_LIVE" not in text and "ApiTestCase" not in text and "restart" not in text.lower() + settings = [source for source in sources if 'PROJECT = "" # @param {type:"string"}' in source] + assert len(settings) == 1 + namespace = {} + with pytest.raises(Exception) as stopped: + exec(compile(settings[0], "settings", "exec"), namespace) + assert type(stopped.value).__name__ == "Stop" + assert stopped.value._render_traceback_() == [f"Stopped: {stopped.value}"] + assert len(_cells("01_quickstart.ipynb")) <= 10 + + +def test_live_steps_in_03_check_the_switches_and_the_project(): + cells = _code("03_bring_your_own_cnn.ipynb") settings = [cell for cell in cells if re.search(r"^RUN_LIVE = False\b", _source(cell), re.M)] - assert len(settings) == 1, "one settings cell with RUN_LIVE = False" - live = [cell for cell in cells if "live" in cell["metadata"].get("tags", [])] - if live: - assert f'PROJECT = "{PLACEHOLDER_PROJECT}"' in _source(settings[0]) + assert len(settings) == 1 and f'PROJECT = "{PLACEHOLDER_PROJECT}"' in _source(settings[0]) for cell in cells: - source = _source(cell) - if cell in live: - assert "live_step(" in source, "a live cell checks RUN_LIVE and the project first" - else: - assert not LIVE_CALLS.search(source), source[:200] + if _live(cell): + assert "live_step(" in _source(cell) @pytest.mark.parametrize("name", sorted(EXPECTED)) @@ -101,17 +133,18 @@ def test_no_private_values(name): text = "\n".join(_source(cell) for cell in _cells(name)) assert not EMAIL.search(text) assert not HOME.search(text) + assert "billed" not in text.lower() public = _public_assets() for match in ASSET.finditer(text): project = match.group(1) # The placeholder, the settings variable, or a public weight table. assert project in (PLACEHOLDER_PROJECT, "{PROJECT}", "PROJECT") or match.group(0) in public, match.group(0) for value in PROJECT_VALUE.findall(text): - assert value == PLACEHOLDER_PROJECT, value + assert value in (PLACEHOLDER_PROJECT, ""), value @contextlib.contextmanager -def _offline(): +def _offline(install_signatures=False): """Refuse every Earth Engine data call and authentication; restore the client afterwards.""" import ee from demos.code_editor.build import CloudCall, _REFUSED_DATA_CALLS @@ -123,39 +156,57 @@ def call(*args, **kwargs): install = ee.data._install_cloud_api_resource limit, path = sys.getrecursionlimit(), list(sys.path) + case = None try: + if install_signatures: + from ee.apitestcase import ApiTestCase + case = ApiTestCase() + case.setUp() with contextlib.ExitStack() as stack: for name in _REFUSED_DATA_CALLS: if hasattr(ee.data, name): stack.enter_context(mock.patch.object(ee.data, name, refuse(name))) - # ee.Initialize stays: the notebooks' ApiTestCase calls it with its offline client. - # The static check above keeps it out of every cell that is not tagged live. + # ee.Initialize stays: ApiTestCase calls it with its offline client. The static + # check above keeps it out of every cell that is not tagged live. stack.enter_context(mock.patch.object(ee, "Authenticate", refuse("Authenticate"))) yield finally: + if case is not None: + case.tearDown() ee.data._install_cloud_api_resource = install ee.Reset() sys.setrecursionlimit(limit) sys.path[:] = path -def _run(name, monkeypatch): +def _exec(name, index, cell, namespace): + lines = [line for line in _source(cell).split("\n") if not line.lstrip().startswith(("%", "!"))] + exec(compile("\n".join(lines), f"{name}[{index}]", "exec"), namespace) + + +def _run(name, monkeypatch, *, after_settings=None, signatures=False, then=None): for module in EXPECTED[name]: pytest.importorskip(module) import torch monkeypatch.chdir(ROOT) namespace = {"__name__": "__notebook__"} threads = torch.get_num_threads() - with _offline(): + with _offline(signatures): try: for index, cell in enumerate(_cells(name)): - if cell["cell_type"] != "code" or "live" in cell["metadata"].get("tags", []): + if cell["cell_type"] != "code" or _live(cell): continue - lines = [line for line in _source(cell).split("\n") - if not line.lstrip().startswith(("%", "!"))] - exec(compile("\n".join(lines), f"{name}[{index}]", "exec"), namespace) + try: + _exec(name, index, cell, namespace) + except Exception as error: + if type(error).__name__ != "Stop" or after_settings is None: + raise + if after_settings and "@param" in _source(cell): + namespace.update(after_settings) if "RUN_LIVE" in namespace: assert namespace["RUN_LIVE"] is False + if then: + then(namespace) finally: if "offline" in namespace: namespace["offline"].tearDown() @@ -165,12 +216,23 @@ def _run(name, monkeypatch): def test_quickstart_runs_offline_and_builds_the_recorded_graph(monkeypatch): import ee - namespace = _run("01_quickstart.ipynb", monkeypatch) - assert namespace["OUTPUT_BANDS"] == ["logit_0", "logit_1", "water_probability", "water_mask", "coverage"] - assert namespace["scene_id"] == namespace["PRESET"]["scene"] + from demos.code_editor import build + place = {"crs": PRESET["crs"], "x": PRESET["transform"][2], "y": PRESET["transform"][5]} + record = json.loads((ROOT / "docs/evidence/s03-template-identity.json").read_text()) + preset = next(aoi for aoi in record["aois"] if aoi["name"] == "segwater_sun_moon_lake_1x1") + found = {} + + def identity(namespace): + found["sha256"] = build.sha256(namespace["prediction"].serialize()) + + namespace = _run("01_quickstart.ipynb", monkeypatch, signatures=True, then=identity, + after_settings={"PROJECT": "offline-check", "place": place, + "scene_id": PRESET["scene_id"]}) + assert namespace["BANDS"] == ["logit_0", "logit_1", "water_probability", "water_mask", "coverage"] + assert namespace["ASSET_NAME"].startswith("tensor_ee_quickstart_") if ee.__version__ == namespace["METADATA"]["earthengineApiVersion"]: # predict_segwater's graph for the preset chip, built with the trained checkpoint. - assert namespace["graph_sha256"] == namespace["preset"]["library_sha256"] + assert found["sha256"] == preset["library_sha256"] def test_is_my_model_supported_runs_offline(monkeypatch): @@ -179,8 +241,23 @@ def test_is_my_model_supported_runs_offline(monkeypatch): assert namespace["b0"].report.lines()[0].startswith("segformer_b0: 20 nodes; 18 with a validated record, 18 inside") -def test_bring_your_own_cnn_runs_offline(monkeypatch): - namespace = _run("03_bring_your_own_cnn.ipynb", monkeypatch) +def test_bring_your_own_cnn_runs_offline_and_goes_live_without_a_restart(monkeypatch): + import ee + + def switch_to_live(namespace): + # Rerunning the connect cell with RUN_LIVE = True undoes the offline signatures. + connect = next((index, cell) for index, cell in enumerate(_cells("03_bring_your_own_cnn.ipynb")) + if "ApiTestCase()" in _source(cell)) + namespace["RUN_LIVE"] = True + offline = namespace["offline"] + _exec("03_bring_your_own_cnn.ipynb", *connect, namespace) + assert "offline" not in namespace + assert ee.data._install_cloud_api_resource is namespace["_install_cloud_api_resource"] + assert ee.data.getAlgorithms is offline.old_get_algorithms + assert ee.data.getMapId is offline.old_get_map_id + namespace["RUN_LIVE"] = False + + namespace = _run("03_bring_your_own_cnn.ipynb", monkeypatch, then=switch_to_live) assert namespace["CHECKPOINT"] is None assert (namespace["halo"], namespace["alignment"]) == (112, 8) assert len(namespace["region_plan"].tiles) == 2 and len(namespace["test_plan"].tiles) == 1 From 8411814f349cdce0c47ca598896b97d7af434783 Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:18:41 +0800 Subject: [PATCH 3/9] Quickstart: a plain ready line, export progress, and a map that waits and loads in Colab Step 5 says the model is set up for the chip and scene instead of listing bands. Step 6 links the Earth Engine Tasks page and prints a short status every minute while it waits. Step 7 waits for a running export, stops with one line if step 6 has not run, centres on the chip, and uses CartoDB Positron and Esri World Imagery as base maps, since OpenStreetMap's tiles refuse Colab. --- notebooks/01_quickstart.ipynb | 47 +++++++++++++++++++++++++++++------ tests/test_notebooks.py | 25 ++++++++++++++++++- 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/notebooks/01_quickstart.ipynb b/notebooks/01_quickstart.ipynb index 7418b40..51d965d 100644 --- a/notebooks/01_quickstart.ipynb +++ b/notebooks/01_quickstart.ipynb @@ -67,6 +67,22 @@ " return [f\"Stopped: {self}\"]\n", "\n", "\n", + "TASKS_URL = \"https://code.earthengine.google.com/tasks\"\n", + "\n", + "\n", + "def wait_for(task, poll=20):\n", + " \"\"\"Wait for an export, with a short status line every minute; stop if it fails.\"\"\"\n", + " started, polls = time.time(), 0\n", + " while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", + " time.sleep(poll)\n", + " polls += 1\n", + " if polls % 3 == 0:\n", + " minutes = (time.time() - started) / 60\n", + " print(f\" still running, {minutes:.0f} min ({state.lower()})...\", flush=True)\n", + " if state != \"COMPLETED\":\n", + " raise Stop(f\"the export ended {state}: {task.status().get('error_message', '')} (see {TASKS_URL})\")\n", + "\n", + "\n", "if not PROJECT.strip():\n", " raise Stop(\"enter your Cloud project ID in PROJECT above, then run all cells again.\")" ], @@ -167,7 +183,8 @@ "prediction = (ee.ImageCollection.fromImages([chip]).mosaic().select(BANDS)\n", " .setDefaultProjection(grid.projection()).clip(grid.geometry())\n", " .updateMask(grid.mask()).toFloat())\n", - "print(\"Output bands:\", \", \".join(BANDS))" + "print(f\"Ready: the trained U-Net is set up for your chip and scene \"\n", + " f\"({len(METADATA['weightAssets'])} public weight tables). Nothing has run yet; step 6 runs it.\")" ], "execution_count": null, "outputs": [] @@ -183,17 +200,15 @@ "source": [ "# @title 6. Run the model (an export of about 65 EECU-seconds, a few minutes)\n", "# @markdown Earth Engine computes the chip in a batch export and saves it as a new image asset.\n", + "# @markdown Follow it in the [Tasks tab](https://code.earthengine.google.com/tasks) too.\n", "asset_id = f\"projects/{PROJECT}/assets/{ASSET_NAME}\"\n", "task = ee.batch.Export.image.toAsset(\n", " image=prediction.set({\"source_scene\": scene_id}), description=\"tensor_ee_quickstart\",\n", " assetId=asset_id, region=grid.geometry(), crs=grid.crs, crsTransform=list(grid.transform),\n", " maxPixels=224 * 224 * len(BANDS), pyramidingPolicy={\".default\": \"sample\"})\n", "task.start()\n", - "print(\"Running; you can also watch it in the Code Editor's Tasks tab.\")\n", - "while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", - " time.sleep(20)\n", - "if state != \"COMPLETED\":\n", - " raise Stop(f\"the export ended {state}: {task.status().get('error_message', '')}\")\n", + "print(\"Running; follow it here or at\", TASKS_URL)\n", + "wait_for(task)\n", "print(\"Done:\", asset_id)" ], "execution_count": null, @@ -209,9 +224,20 @@ }, "source": [ "# @title 7. See the radar image and the water map\n", - "# @markdown Use the layer control (top right) to switch layers.\n", + "# @markdown Waits for step 6's export if it is still running. Layers: radar, water probability,\n", + "# @markdown water mask (1 where the probability is at least 0.5); switch them top right.\n", "import folium\n", "\n", + "if \"task\" in globals() and task.status()[\"state\"] != \"COMPLETED\":\n", + " print(\"Waiting for step 6's export to finish...\")\n", + " wait_for(task)\n", + "elif \"task\" not in globals():\n", + " try:\n", + " ee.data.getAsset(f\"projects/{PROJECT}/assets/{ASSET_NAME}\")\n", + " except ee.EEException:\n", + " raise Stop(\"run step 6 first: it creates the water map this step shows.\") from None\n", + "asset_id = f\"projects/{PROJECT}/assets/{ASSET_NAME}\"\n", + "\n", "def layer(image, visualization, name, show=True):\n", " url = ee.Image(image).getMapId(visualization)[\"tile_fetcher\"].url_format\n", " return folium.raster_layers.TileLayer(\n", @@ -219,7 +245,12 @@ " attr=\"Google Earth Engine; contains modified Copernicus Sentinel data\")\n", "\n", "result = ee.Image(asset_id)\n", - "water_map = folium.Map(location=[LATITUDE, LONGITUDE], zoom_start=14)\n", + "centre_lon, centre_lat = grid.geometry().centroid(1).transform(\"EPSG:4326\", 1).coordinates().getInfo()\n", + "water_map = folium.Map(location=[centre_lat, centre_lon], zoom_start=14, tiles=\"CartoDB positron\")\n", + "folium.raster_layers.TileLayer(\n", + " tiles=\"https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}\",\n", + " attr=\"Tiles © Esri — Source: Esri, Maxar, Earthstar Geographics, and the GIS User Community\",\n", + " name=\"Esri World Imagery\").add_to(water_map)\n", "layer(ee.Image(scene_id).select(\"VV\"), {\"min\": -25, \"max\": 0}, \"Sentinel-1 VV (dB)\").add_to(water_map)\n", "layer(result.select(\"water_probability\"), {\"min\": 0, \"max\": 1, \"palette\": [\"ffffff\", \"08306b\"]},\n", " \"Water probability\").add_to(water_map)\n", diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py index 7e99427..1bd17dc 100644 --- a/tests/test_notebooks.py +++ b/tests/test_notebooks.py @@ -37,7 +37,7 @@ COLAB = "https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/" # Calls that reach Earth Engine or start work: only cells tagged "live" may make them. LIVE_CALLS = re.compile(r"ee\.Initialize|ee\.Authenticate|\.getInfo\(|\.start\(\)|computePixels|" - r"getMapId|\.status\(\)|export_weights\(") + r"getMapId|export_weights\(") # task.status() needs a started task EMAIL = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+\.[A-Za-z0-9.-]+") HOME = re.compile(r"/Users/|/home/|[A-Za-z]:\\\\Users|~/") ASSET = re.compile(r"projects/([^/\s\"'`]+)/assets/[^\s\"'`)]*") @@ -119,6 +119,29 @@ def test_the_quickstart_has_one_setting_and_stops_cleanly_without_it(): assert len(_cells("01_quickstart.ipynb")) <= 10 +def test_the_quickstart_waits_with_progress_and_stops_on_a_failed_export(capsys): + settings = next(_source(cell) for cell in _code("01_quickstart.ipynb") if "def wait_for" in _source(cell)) + namespace = {} + with pytest.raises(Exception): + exec(compile(settings, "settings", "exec"), namespace) # stops: PROJECT is empty + + class Task: + def __init__(self, states): + self.states = list(states) + + def status(self): + state = self.states.pop(0) if len(self.states) > 1 else self.states[0] + return {"state": state, "error_message": "out of memory"} + + with mock.patch.object(namespace["time"], "sleep", lambda seconds: None): + namespace["wait_for"](Task(["READY"] + ["RUNNING"] * 6 + ["COMPLETED"])) + assert capsys.readouterr().out.count("still running") == 2 + with pytest.raises(Exception) as stopped: + namespace["wait_for"](Task(["RUNNING", "FAILED"])) + assert type(stopped.value).__name__ == "Stop" and "out of memory" in str(stopped.value) + assert namespace["TASKS_URL"] in str(stopped.value) + + def test_live_steps_in_03_check_the_switches_and_the_project(): cells = _code("03_bring_your_own_cnn.ipynb") settings = [cell for cell in cells if re.search(r"^RUN_LIVE = False\b", _source(cell), re.M)] From 76019fd35f586509c5c08435711c8bc376903b87 Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:20:28 +0800 Subject: [PATCH 4/9] Quickstart map: Esri base layers, which need no API key --- notebooks/01_quickstart.ipynb | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/notebooks/01_quickstart.ipynb b/notebooks/01_quickstart.ipynb index 51d965d..250471d 100644 --- a/notebooks/01_quickstart.ipynb +++ b/notebooks/01_quickstart.ipynb @@ -246,9 +246,12 @@ "\n", "result = ee.Image(asset_id)\n", "centre_lon, centre_lat = grid.geometry().centroid(1).transform(\"EPSG:4326\", 1).coordinates().getInfo()\n", - "water_map = folium.Map(location=[centre_lat, centre_lon], zoom_start=14, tiles=\"CartoDB positron\")\n", + "ESRI = \"https://server.arcgisonline.com/ArcGIS/rest/services/{}/MapServer/tile/{{z}}/{{y}}/{{x}}\"\n", + "water_map = folium.Map(location=[centre_lat, centre_lon], zoom_start=14, tiles=None)\n", + "folium.raster_layers.TileLayer(tiles=ESRI.format(\"Canvas/World_Light_Gray_Base\"),\n", + " attr=\"Tiles © Esri\", name=\"Esri Light Gray\").add_to(water_map)\n", "folium.raster_layers.TileLayer(\n", - " tiles=\"https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}\",\n", + " tiles=ESRI.format(\"World_Imagery\"),\n", " attr=\"Tiles © Esri — Source: Esri, Maxar, Earthstar Geographics, and the GIS User Community\",\n", " name=\"Esri World Imagery\").add_to(water_map)\n", "layer(ee.Image(scene_id).select(\"VV\"), {\"min\": -25, \"max\": 0}, \"Sentinel-1 VV (dB)\").add_to(water_map)\n", From ad86f658a77fa9691a4d8c4411cc3188479f3e71 Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:35:26 +0800 Subject: [PATCH 5/9] Notebook 02: install from the repository, and trim the evidence notes --- notebooks/02_is_my_model_supported.ipynb | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/notebooks/02_is_my_model_supported.ipynb b/notebooks/02_is_my_model_supported.ipynb index d6ee1fd..c362a99 100644 --- a/notebooks/02_is_my_model_supported.ipynb +++ b/notebooks/02_is_my_model_supported.ipynb @@ -36,10 +36,8 @@ "id": "5458608d", "metadata": {}, "source": [ - "%pip install --quiet \"tensor-ee[swin]\"\n", - "# Grouped and depthwise convolutions (ResNeXt and MobileNet encoders, for example) became\n", - "# supported after the 0.1.0 release. To check such a model, install from the repository instead:\n", - "# %pip install --quiet \"tensor-ee[swin] @ git+https://github.com/noelEOS/tensor-ee\"" + "# Install from the repository: support for more layers is added there between releases.\n", + "%pip install --quiet \"tensor-ee[swin] @ git+https://github.com/noelEOS/tensor-ee\"" ], "execution_count": null, "outputs": [] @@ -143,11 +141,6 @@ "metadata": {}, "source": [ "- **\"Supported\"** means TensorEE can express the model; it is not yet a claim about the numbers.\n", - "- **Evidence so far:** the blocks agree with PyTorch at their validated settings, and one trained\n", - " ResNet-18 U-Net prepared this way matched PyTorch in Earth Engine to 1.2e-5 (64 × 64 input).\n", - " Other encoders have not been run live as whole models.\n", - "- **Checked offline the same way:** ResNet-18, 34, 50 and VGG-11 (BatchNorm) encoders in `Unet`\n", - " and `UnetPlusPlus`.\n", "- Supported is not accurate: accuracy depends on training." ] }, @@ -160,7 +153,7 @@ "\n", "### A LinkNet: transposed convolutions\n", "\n", - "- LinkNet's decoder upsamples with `ConvTranspose2d`, which has no CNN building block." + "- LinkNet's decoder upsamples with `ConvTranspose2d`, which has no CNN building block yet. Support for more layers is being added, so install from the repository (step 1) to get the latest." ] }, { From eadbbb77c36ff6840a9a7323a215747efac75b8d Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:48:30 +0800 Subject: [PATCH 6/9] Notebook 03: install from the repository --- notebooks/03_bring_your_own_cnn.ipynb | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/notebooks/03_bring_your_own_cnn.ipynb b/notebooks/03_bring_your_own_cnn.ipynb index 851799d..27c88fd 100644 --- a/notebooks/03_bring_your_own_cnn.ipynb +++ b/notebooks/03_bring_your_own_cnn.ipynb @@ -35,7 +35,8 @@ "metadata": {}, "source": [ "# @title 1. Install\n", - "%pip install --quiet \"tensor-ee[segwater]\"" + "# Install from the repository: support for more layers is added there between releases.\n", + "%pip install --quiet \"tensor-ee[segwater] @ git+https://github.com/noelEOS/tensor-ee\"" ], "execution_count": null, "outputs": [] From 9769a02bbfe8c04ed30748335879f1d9a37da936 Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:15:02 +0800 Subject: [PATCH 7/9] Reader notebooks: plain-language settings, a checkbox per Earth Engine step, and a message from every cell - 03: settings grouped, each line says whether to change it; USE_EARTH_ENGINE replaces RUN_LIVE; the four Earth Engine steps each get a checkbox in their own cell, run one at a time; the helper code moves into the titled connect cell; every cell prints what happened or why not; steps renumbered so titles, switches and text agree. - 02: setup in one titled cell that says it is ready; terms explained where they first appear. - 01: wording pass; install and settings cells now confirm. - start-here: no coding agent needed; plainer wording of the checks. - tests: follow the new 03 names, and check its four checkboxes. --- docs/start-here.md | 16 +- notebooks/01_quickstart.ipynb | 22 +- notebooks/02_is_my_model_supported.ipynb | 342 +++++++-------- notebooks/03_bring_your_own_cnn.ipynb | 533 +++++++++++++---------- tests/test_notebooks.py | 20 +- 5 files changed, 491 insertions(+), 442 deletions(-) diff --git a/docs/start-here.md b/docs/start-here.md index 930dc4c..b649611 100644 --- a/docs/start-here.md +++ b/docs/start-here.md @@ -1,7 +1,10 @@ # Start here -For anyone with an Earth Engine account and some Python who wants to try TensorEE. A browser -is enough: the notebooks run in Google Colab. +For anyone with an Earth Engine account and some Python who wants to try TensorEE. + +- **A browser is enough:** the notebooks run in Google Colab. +- **No coding agent needed:** TensorEE is an ordinary Python library. You run the notebooks + cell by cell, as you would any other. ## What TensorEE does, and does not do @@ -31,8 +34,10 @@ is enough: the notebooks run in Google Colab. refused and why. - [Notebook 03, bring your own CNN](../notebooks/03_bring_your_own_cnn.ipynb) ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)): - a trained segmentation-models-pytorch U-Net from checkpoint to an exported region. It runs - offline first; each Earth Engine step has its own switch. + a trained segmentation-models-pytorch U-Net from checkpoint to an exported region. Try it + offline with a demo model first; then each Earth Engine step (upload the weights, export a + test tile, compare it with PyTorch, export the region) has its own checkbox, run one at a + time. For details, see [Getting started](getting-started.md) and the [guides](README.md). @@ -64,7 +69,8 @@ For details, see [Getting started](getting-started.md) and the [guides](README.m ## What has been checked -- **Graph identity:** the Earth Engine computation is byte for byte one checked before. +- **Same operations every time:** the Earth Engine operations TensorEE builds for its reference + models are compared byte for byte with saved copies, so an unintended change is caught. - **Offline agreement:** each building block, computed in NumPy, matches PyTorch. - **Live agreement:** Earth Engine's result matches PyTorch on the same input, within 1e-4 in the logits. For the trained U-Net's recorded chip: at most 1.35e-5 diff --git a/notebooks/01_quickstart.ipynb b/notebooks/01_quickstart.ipynb index 250471d..97a88ce 100644 --- a/notebooks/01_quickstart.ipynb +++ b/notebooks/01_quickstart.ipynb @@ -36,7 +36,8 @@ "\n", "REPOSITORY = next(folder.resolve() for folder in (Path(\"tensor-ee\"), Path(\".\"), Path(\"..\"))\n", " if (folder / \"demos\" / \"code_editor\" / \"segwater.js\").is_file())\n", - "sys.path.insert(0, str(REPOSITORY))" + "sys.path.insert(0, str(REPOSITORY))\n", + "print(\"TensorEE is installed.\")" ], "execution_count": null, "outputs": [] @@ -84,7 +85,9 @@ "\n", "\n", "if not PROJECT.strip():\n", - " raise Stop(\"enter your Cloud project ID in PROJECT above, then run all cells again.\")" + " raise Stop(\"enter your Cloud project ID in PROJECT above, then run all cells again.\")\n", + "print(f\"Settings saved: project {PROJECT}, point ({LONGITUDE}, {LATITUDE}), \"\n", + " f\"scenes from {START_DATE} to {END_DATE}.\")" ], "execution_count": null, "outputs": [] @@ -104,7 +107,7 @@ "\n", "ee.Authenticate()\n", "ee.Initialize(project=PROJECT)\n", - "print(\"Signed in to\", PROJECT)" + "print(f\"Earth Engine is ready — signed in to {PROJECT}.\")" ], "execution_count": null, "outputs": [] @@ -155,7 +158,7 @@ " raise Stop(f\"no Sentinel-1 scene covers the whole chip between {START_DATE} and {END_DATE}; \"\n", " \"widen the dates or move the point.\")\n", "scene_id = complete[0]\n", - "print(\"Chip:\", crs, (place[\"x\"], place[\"y\"]), \"| Scene:\", scene_id)" + "print(f\"Found the chip ({CHIP} × {CHIP} px in {crs}) and a Sentinel-1 scene that covers it: {scene_id}\")" ], "execution_count": null, "outputs": [] @@ -165,13 +168,14 @@ "id": "c035f536", "metadata": {}, "source": [ - "# @title 5. Build the model's computation\n", - "# @markdown Fills the trained U-Net's stored computation with your chip and scene. Nothing runs yet.\n", + "# @title 5. Set up the model for your chip\n", + "# @markdown Loads the trained U-Net, stored in the repository as Earth Engine operations, and points\n", + "# @markdown it at your chip and scene. Nothing runs yet.\n", "import ee\n", "from demos.code_editor import build as segwater_module\n", "from tensor_ee import Grid\n", "\n", - "sys.setrecursionlimit(max(sys.getrecursionlimit(), 20000)) # the network is a deep expression\n", + "sys.setrecursionlimit(max(sys.getrecursionlimit(), 20000)) # the network is deeply nested\n", "TEMPLATE, METADATA = segwater_module.parse_module(\n", " (REPOSITORY / \"demos\" / \"code_editor\" / \"segwater.js\").read_text())\n", "BANDS = list(segwater_module.OUTPUT_BANDS)\n", @@ -207,9 +211,9 @@ " assetId=asset_id, region=grid.geometry(), crs=grid.crs, crsTransform=list(grid.transform),\n", " maxPixels=224 * 224 * len(BANDS), pyramidingPolicy={\".default\": \"sample\"})\n", "task.start()\n", - "print(\"Running; follow it here or at\", TASKS_URL)\n", + "print(\"Export started; it usually takes a few minutes. Follow it here or at\", TASKS_URL)\n", "wait_for(task)\n", - "print(\"Done:\", asset_id)" + "print(\"Done: the water map is saved as\", asset_id)" ], "execution_count": null, "outputs": [] diff --git a/notebooks/02_is_my_model_supported.ipynb b/notebooks/02_is_my_model_supported.ipynb index c362a99..696f4df 100644 --- a/notebooks/02_is_my_model_supported.ipynb +++ b/notebooks/02_is_my_model_supported.ipynb @@ -2,81 +2,63 @@ "cells": [ { "cell_type": "markdown", - "id": "5df6cfd8", + "id": "03e548c9", "metadata": {}, "source": [ "# Is my model supported?\n", "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/02_is_my_model_supported.ipynb)\n", "\n", - "TensorEE runs a model only if every operation in it has a building block it can compute\n", - "exactly. This notebook checks a PyTorch model and says what is supported, what is refused, and\n", - "what to do next.\n", + "TensorEE runs a model only if it has an exact Earth Engine version (a *building block*) of\n", + "every layer in it. This notebook checks a PyTorch model and tells you what is supported, what\n", + "is refused, and what to do next.\n", "\n", - "- **Fully offline:** no Earth Engine account, no requests, no downloads. The examples use\n", - " random weights, which is all a support check needs: it reads the model's structure, not its\n", - " predictions.\n", + "- **Fully offline:** no Earth Engine account, no requests, no downloads.\n", + "- **Random weights are enough:** the check reads the model's layers, not its predictions.\n", "- **To run:** **Runtime → Run all**. Then put your own model in step 6.\n", - "- **Examples:** a supported U-Net; two refused models; an attention model through the converter." - ] - }, - { - "cell_type": "markdown", - "id": "a8601d7f", - "metadata": {}, - "source": [ - "## 1. Install\n", - "\n", - "- The `swin` extra brings segmentation-models-pytorch 0.5.0 and timm 1.0.30, the versions the\n", - " library was checked with." + "- **Examples:** a supported U-Net, two refused models, and attention models such as SegFormer." ] }, { "cell_type": "code", - "id": "5458608d", - "metadata": {}, - "source": [ - "# Install from the repository: support for more layers is added there between releases.\n", - "%pip install --quiet \"tensor-ee[swin] @ git+https://github.com/noelEOS/tensor-ee\"" - ], "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "c424214b", + "id": "ec064c72", "metadata": {}, + "outputs": [], "source": [ - "## 2. Setup\n", - "\n", - "- The chip size the examples are checked at, and a helper that explains a report." + "# @title 1. Install\n", + "# Installs from the repository, which gets support for more layers between releases. The\n", + "# \"swin\" extra brings segmentation-models-pytorch 0.5.0 and timm 1.0.30, the versions TensorEE\n", + "# was checked with.\n", + "%pip install --quiet \"tensor-ee[swin] @ git+https://github.com/noelEOS/tensor-ee\"\n", + "print(\"TensorEE is installed.\")" ] }, { "cell_type": "code", - "id": "295d1ada", - "metadata": {}, - "source": [ - "INPUT_SHAPE = (1, 2, 64, 64) # one chip: (1, bands, height, width)" - ], "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "cf1e5641", - "metadata": {}, + "id": "2166edea", + "metadata": { + "cellView": "form" + }, + "outputs": [], "source": [ + "# @title 2. Set up the checker\n", + "# @markdown Loads the checker, and a helper that explains its report in plain words.\n", "from collections import Counter\n", "\n", "import torch\n", "import segmentation_models_pytorch as smp\n", "\n", "from tensor_ee import check_compatibility\n", + "from tensor_ee.converter.convert import convert\n", + "from tensor_ee.core.errors import UnsupportedOperation\n", + "from tensor_ee.core.grid import Grid\n", "from tensor_ee.smp import adapt_smp_model\n", "\n", "torch.set_num_threads(2)\n", "torch.manual_seed(0) # random weights, the same every run\n", + "INPUT_SHAPE = (1, 2, 64, 64) # the examples' input: (1, bands, height, width)\n", "\n", "\n", "def explain(report):\n", @@ -84,10 +66,10 @@ " facts = report.to_dict()\n", " if facts[\"status\"] == \"supported\":\n", " storage = facts[\"weight_storage\"]\n", - " print(\"SUPPORTED: every operation in this model has a TensorEE building block.\")\n", - " print(f\" {facts['parameter_count']:,} parameters. After BatchNorm is folded into the \"\n", + " print(\"SUPPORTED: TensorEE has a building block for every layer in this model.\")\n", + " print(f\" {facts['parameter_count']:,} parameters. After BatchNorm is merged into the \"\n", " f\"convolutions, {storage['coefficient_count']:,} numbers would be uploaded as weight assets.\")\n", - " print(f\" Output for one chip: {facts['output_shape']} (batch, classes, height, width).\")\n", + " print(f\" Output for one input: {facts['output_shape']} (batch, classes, height, width).\")\n", " return\n", " print(f\"NOT SUPPORTED (status: {facts['status']}).\")\n", " if not facts[\"evaluation\"][\"model_eval\"]:\n", @@ -95,225 +77,264 @@ " if not facts[\"evaluation\"][\"float32_parameters_and_buffers\"]:\n", " print(\" Some weights are not float32: TensorEE computes in float32 at most.\")\n", " if facts[\"trace\"][\"status\"] != \"success\":\n", - " print(\" The model's code could not be traced:\", facts[\"trace\"][\"error\"])\n", + " print(\" The model's code could not be read:\", facts[\"trace\"][\"error\"])\n", " reasons = Counter((blocker[\"code\"], blocker[\"message\"]) for blocker in facts[\"blockers\"])\n", " for (code, message), count in reasons.most_common(8):\n", " print(f\" {count} x {message} [{code}]\")\n", " if len(reasons) > 8:\n", " print(f\" ... and {len(reasons) - 8} more kinds of refusal (report.to_dict()['blockers'] lists them all)\")\n", " if facts[\"unknowns\"]:\n", - " print(f\" {len(facts['unknowns'])} later operations were not checked, because they depend \"\n", - " \"on the refused ones.\")" - ], - "execution_count": null, - "outputs": [] + " print(f\" {len(facts['unknowns'])} later layers were not checked, because they depend \"\n", + " \"on the refused ones.\")\n", + "\n", + "\n", + "def check_my_model(model, input_shape):\n", + " \"\"\"Try the CNN check, then, if that does not support the model, the converter.\"\"\"\n", + " model = model.eval()\n", + " print(\"CNN check (U-Nets and other plain CNNs):\")\n", + " try:\n", + " if type(model).__module__.startswith(\"segmentation_models_pytorch\"):\n", + " model_for_check = adapt_smp_model(model, input_shape=input_shape)\n", + " else:\n", + " model_for_check = model # a plain PyTorch CNN is checked as it is\n", + " report = check_compatibility(model_for_check, input_shape)\n", + " explain(report)\n", + " if report.status == \"supported\":\n", + " return report\n", + " except (ValueError, RuntimeError, UnsupportedOperation) as refusal:\n", + " print(\"NOT SUPPORTED:\", refusal)\n", + "\n", + " print()\n", + " print(\"Converter (attention models):\")\n", + " height, width = input_shape[2], input_shape[3]\n", + " try:\n", + " try:\n", + " result = convert(model, input_shape, Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000),\n", + " width, height), name=\"my_model\")\n", + " except UnsupportedOperation as refusal:\n", + " if \"source_grid=None\" not in str(refusal):\n", + " raise\n", + " result = convert(model, input_shape, name=\"my_model\") # a Swin encoder starts on its own grid\n", + " print(result.report.lines()[0])\n", + " print(\"Every block's verdict: print(MY_RESULT.report.text())\")\n", + " return result\n", + " except UnsupportedOperation as refusal:\n", + " print(\"NOT SUPPORTED:\", refusal.code, \"-\", str(refusal)[:600])\n", + " except Exception as failure: # for example, torch.export could not read the model\n", + " print(\"NOT SUPPORTED: the converter could not read this model:\", type(failure).__name__, failure)\n", + "\n", + "\n", + "print(\"The checker is ready.\")" + ] }, { "cell_type": "markdown", - "id": "00d29019", + "id": "af9bb1a2", "metadata": {}, "source": [ "## 3. A supported model: a U-Net with a ResNet-18 encoder\n", "\n", - "- The architecture of notebook 01's water model, with random weights.\n", - "- `adapt_smp_model` prepares an SMP model at a fixed chip size, and checks that it gives exactly\n", - " the original's numbers on a test input.\n", - "- `check_compatibility` checks every operation." + "- The same architecture as notebook 01's water model, with random weights.\n", + "- `adapt_smp_model` prepares a copy of an SMP model for one input size, and checks that the\n", + " copy gives exactly the original's output on a test input.\n", + "- `check_compatibility` then checks every layer." ] }, { "cell_type": "code", - "id": "63b49161", + "execution_count": null, + "id": "9498d2f1", "metadata": {}, + "outputs": [], "source": [ "unet = smp.Unet(\"resnet18\", encoder_weights=None, in_channels=2, classes=2).eval()\n", "prepared = adapt_smp_model(unet, input_shape=INPUT_SHAPE)\n", - "print(\"Prepared model matches the original on a test input:\",\n", - " prepared.tensor_ee_smp_metadata[\"torch_equal\"])\n", + "if prepared.tensor_ee_smp_metadata[\"torch_equal\"]:\n", + " print(\"The prepared copy gives exactly the original's output on a test input.\")\n", "explain(check_compatibility(prepared, INPUT_SHAPE))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "443e3ab9", + "id": "65fcf0bc", "metadata": {}, "source": [ - "- **\"Supported\"** means TensorEE can express the model; it is not yet a claim about the numbers.\n", + "- **\"Supported\"** means TensorEE can build the model from its blocks. That is not yet a claim\n", + " about the numbers: whether Earth Engine gives the same as PyTorch is checked on a real run\n", + " (notebook 03, step 10).\n", "- Supported is not accurate: accuracy depends on training." ] }, { "cell_type": "markdown", - "id": "0692e745", + "id": "af182bec", "metadata": {}, "source": [ "## 4. Refused models\n", "\n", "### A LinkNet: transposed convolutions\n", "\n", - "- LinkNet's decoder upsamples with `ConvTranspose2d`, which has no CNN building block yet. Support for more layers is being added, so install from the repository (step 1) to get the latest." + "- LinkNet's decoder upsamples with `ConvTranspose2d`, which has no building block yet. Support\n", + " for more layers is being added; installing from the repository (step 1) gets the latest." ] }, { "cell_type": "code", - "id": "944ada38", + "execution_count": null, + "id": "26e676de", "metadata": {}, + "outputs": [], "source": [ "linknet = smp.Linknet(\"resnet18\", encoder_weights=None, in_channels=2, classes=2).eval()\n", "explain(check_compatibility(adapt_smp_model(linknet, input_shape=INPUT_SHAPE), INPUT_SHAPE))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "d07d1603", + "id": "f593c0f3", "metadata": {}, "source": [ "### DeepLabV3+: dilated convolutions, and more\n", "\n", - "- Its encoder uses dilated convolutions; the CNN path supports stride 1 or 2, no dilation,\n", - " \"same\" zero padding.\n", + "- Its encoder uses dilated convolutions. TensorEE's convolutions support stride 1 or 2, no\n", + " dilation, and \"same\" zero padding.\n", "- The report also names its adaptive average pooling and bilinear resize." ] }, { "cell_type": "code", - "id": "f1d487d8", + "execution_count": null, + "id": "ed587d4a", "metadata": {}, + "outputs": [], "source": [ "deeplab = smp.DeepLabV3Plus(\"resnet18\", encoder_weights=None, in_channels=2, classes=2).eval()\n", "explain(check_compatibility(adapt_smp_model(deeplab, input_shape=INPUT_SHAPE), INPUT_SHAPE))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "f73d6c10", + "id": "975e2f88", "metadata": {}, "source": [ - "**A refusal is firm:** TensorEE never approximates. It names the blocking operations; the ones\n", - "after them are not checked. Options:\n", + "**A refusal is firm:** TensorEE never approximates. It names the layers it cannot run; the\n", + "layers after them are not checked. Your options:\n", "\n", "- **Change the architecture** (for example a U-Net decoder instead of transposed convolutions)\n", " and retrain.\n", - "- **Ask for the operation:** [open an issue](https://github.com/noelEOS/tensor-ee/issues) with\n", - " the message and your architecture. A new block comes with its own tests and evidence\n", + "- **Ask for the layer:** [open an issue](https://github.com/noelEOS/tensor-ee/issues) with the\n", + " message and your architecture. A new block comes with its own tests and evidence\n", " ([CONTRIBUTING](https://github.com/noelEOS/tensor-ee/blob/main/CONTRIBUTING.md))." ] }, { "cell_type": "markdown", - "id": "b9d42f92", + "id": "25e7a7d4", "metadata": {}, "source": [ "## 5. Attention models: the converter\n", "\n", - "- `adapt_smp_model` refuses transformer models and points to the converter,\n", - " `tensor_ee.converter.convert`.\n", - "- The converter was built for SegFormer-B0 and a Swin-T UPerNet, on 224 × 224 chips." + "- Transformer (attention) models do not go through the CNN check: `adapt_smp_model` refuses\n", + " them and points to a separate tool, the *converter* (`tensor_ee.converter.convert`).\n", + "- The converter was built for two models, SegFormer-B0 and a Swin-T UPerNet, on 224 × 224\n", + " inputs." ] }, { "cell_type": "code", - "id": "35ee28f7", + "execution_count": null, + "id": "515d11e5", "metadata": {}, + "outputs": [], "source": [ - "from tensor_ee.core.errors import UnsupportedOperation\n", - "from tensor_ee.core.grid import Grid\n", - "from tensor_ee.converter.convert import convert\n", - "\n", "segformer = smp.Segformer(\"mit_b0\", encoder_weights=None, in_channels=2, classes=2).eval()\n", "try:\n", " adapt_smp_model(segformer, input_shape=(1, 2, 224, 224))\n", "except ValueError as refusal:\n", " print(\"adapt_smp_model:\", refusal)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "01ed7f1c", + "id": "23fd38cf", "metadata": {}, "source": [ - "- Converting records the model's structure by running its code on placeholder tensors, so\n", - " convert only models whose source you trust.\n", - "- The report gives each block a verdict: **inside** the settings validated in Earth Engine, or\n", - " **outside**." + "- To read a model, the converter runs its code on placeholder inputs, so convert only models\n", + " whose source code you trust.\n", + "- Its report splits the model into steps (*nodes*). A *validated record* is a set of block\n", + " settings already checked in Earth Engine; each block is **inside** its record or **outside**.\n", + "- The first line counts them; the \"layout operations\" only move pixels into and out of the\n", + " model's own layout." ] }, { "cell_type": "code", - "id": "43d0b8d0", + "execution_count": null, + "id": "5121948a", "metadata": {}, + "outputs": [], "source": [ - "chip = Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000), 224, 224)\n", + "chip = Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000), 224, 224) # a placeholder 224 × 224 input area\n", "b0 = convert(segformer, (1, 2, 224, 224), chip, name=\"segformer_b0\")\n", "print(b0.report.lines()[0])" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "4fbb3f2c", + "id": "fffc365e", "metadata": {}, "source": [ - "- SegFormer-B0: all 18 blocks inside; the 2 layout operations only move pixels on and off the\n", - " model's grid.\n", - "- SegFormer-B1 (same structure, wider) converts, but its blocks are **outside**: expressible,\n", - " not yet validated. Outside is reported, not refused." + "- SegFormer-B0: all 18 blocks inside.\n", + "- SegFormer-B1 (the same structure, wider) converts, but most of its blocks are **outside**: TensorEE\n", + " can build them, but they have not been checked in Earth Engine yet. Outside is reported, not\n", + " refused." ] }, { "cell_type": "code", - "id": "8ead7dab", + "execution_count": null, + "id": "a7e06063", "metadata": {}, + "outputs": [], "source": [ "segformer_b1 = smp.Segformer(\"mit_b1\", encoder_weights=None, in_channels=2, classes=2).eval()\n", "b1 = convert(segformer_b1, (1, 2, 224, 224), chip, name=\"segformer_b1\")\n", "print(b1.report.lines()[0])\n", "print(next(node for node in b1.report.nodes if node.id == \"encoder.block1.0\").verdict())" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "82600b1d", + "id": "fb434a3d", "metadata": {}, "source": [ - "- A refusal: SegFormer-B0 at 512 × 512 reduces its keys and values to more than one 7 × 7\n", - " window, which needs a block the library does not have." + "- A refusal: at 512 × 512, SegFormer-B0's attention works on a grid larger than 7 × 7, which\n", + " needs a building block the library does not have yet." ] }, { "cell_type": "code", - "id": "77add088", + "execution_count": null, + "id": "7e25af36", "metadata": {}, + "outputs": [], "source": [ "try:\n", " convert(segformer, (1, 2, 512, 512), Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000), 512, 512))\n", "except UnsupportedOperation as refusal:\n", " print(\"Refused:\", refusal.code)\n", " print(str(refusal)[:400])" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "1b23e018", + "id": "b34a78b5", "metadata": {}, "source": [ "## 6. Your model\n", "\n", - "- Put your model in `MY_MODEL`, built exactly as trained (random weights give the same answer),\n", - " and its chip shape in `MY_INPUT_SHAPE`. For example:\n", + "- Build your model exactly as you trained it (random weights give the same answer), and give\n", + " its input shape. For example:\n", "\n", "```python\n", "MY_MODEL = smp.Unet(\"resnet34\", encoder_weights=None, in_channels=4, classes=3)\n", @@ -321,75 +342,28 @@ "MY_INPUT_SHAPE = (1, 4, 256, 256)\n", "```\n", "\n", - "- The check tries the CNN path (SMP or plain PyTorch CNN), then the converter." + "- The check tries the CNN check first (SMP models and plain PyTorch CNNs), then the converter." ] }, { "cell_type": "code", - "id": "5f2159b6", - "metadata": {}, - "source": [ - "MY_MODEL = None # your torch.nn.Module\n", - "MY_INPUT_SHAPE = (1, 2, 224, 224) # (1, bands, height, width)" - ], "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "382aae7e", + "id": "7c46d932", "metadata": {}, + "outputs": [], "source": [ - "from tensor_ee.core.errors import UnsupportedOperation\n", - "\n", - "\n", - "def check_my_model(model, input_shape):\n", - " \"\"\"Try the CNN path, then, if that does not support the model, the converter.\"\"\"\n", - " model = model.eval()\n", - " print(\"CNN path (adapt_smp_model and check_compatibility):\")\n", - " try:\n", - " if type(model).__module__.startswith(\"segmentation_models_pytorch\"):\n", - " model_for_check = adapt_smp_model(model, input_shape=input_shape)\n", - " else:\n", - " model_for_check = model # a plain PyTorch CNN is checked as it is\n", - " report = check_compatibility(model_for_check, input_shape)\n", - " explain(report)\n", - " if report.status == \"supported\":\n", - " return report\n", - " except (ValueError, RuntimeError, UnsupportedOperation) as refusal:\n", - " print(\"NOT SUPPORTED:\", refusal)\n", - "\n", - " print()\n", - " print(\"Attention-model path (the converter):\")\n", - " height, width = input_shape[2], input_shape[3]\n", - " try:\n", - " try:\n", - " result = convert(model, input_shape, Grid(\"EPSG:3857\", (2.5, 0, 1000, 0, -2.5, 2000),\n", - " width, height), name=\"my_model\")\n", - " except UnsupportedOperation as refusal:\n", - " if \"source_grid=None\" not in str(refusal):\n", - " raise\n", - " result = convert(model, input_shape, name=\"my_model\") # a Swin encoder starts on its own grid\n", - " print(result.report.lines()[0])\n", - " print(\"Every block's verdict: print(MY_RESULT.report.text())\")\n", - " return result\n", - " except UnsupportedOperation as refusal:\n", - " print(\"NOT SUPPORTED:\", refusal.code, \"-\", str(refusal)[:600])\n", - " except Exception as failure: # for example, torch.export could not trace the model\n", - " print(\"NOT SUPPORTED: the converter could not read this model:\", type(failure).__name__, failure)\n", - "\n", + "MY_MODEL = None # change to your torch.nn.Module\n", + "MY_INPUT_SHAPE = (1, 2, 224, 224) # change to (1, bands, height, width)\n", "\n", "if MY_MODEL is None:\n", - " print(\"Put your model in MY_MODEL in the cell above, then run this cell again.\")\n", + " print(\"No model yet: set MY_MODEL above, then run this cell again.\")\n", "else:\n", " MY_RESULT = check_my_model(MY_MODEL, MY_INPUT_SHAPE)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", - "id": "9113a12b", + "id": "9020a2f7", "metadata": {}, "source": [ "## What next\n", @@ -397,13 +371,13 @@ "- **Supported CNN:** [notebook 03](https://github.com/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", " takes it from checkpoint to an exported region.\n", "- **Converted attention model:** the\n", - " [converter guide](https://github.com/noelEOS/tensor-ee/blob/main/docs/converter.md) plans and\n", - " runs its stages.\n", - "- **Outside the validated records:** run it and compare with PyTorch on your input to make that\n", + " [converter guide](https://github.com/noelEOS/tensor-ee/blob/main/docs/converter.md) explains\n", + " how to run it in Earth Engine.\n", + "- **Outside the checked settings:** run it and compare with PyTorch on your input to make that\n", " evidence; the maintainer can advise.\n", "- **Refused:** see step 4. The\n", " [block catalogue](https://github.com/noelEOS/tensor-ee/blob/main/docs/building-blocks.md)\n", - " lists every block and its validated settings." + " lists every block and its checked settings." ] } ], diff --git a/notebooks/03_bring_your_own_cnn.ipynb b/notebooks/03_bring_your_own_cnn.ipynb index 27c88fd..d32a734 100644 --- a/notebooks/03_bring_your_own_cnn.ipynb +++ b/notebooks/03_bring_your_own_cnn.ipynb @@ -2,103 +2,129 @@ "cells": [ { "cell_type": "markdown", - "id": "ab4596ab", + "id": "68ccd5df", "metadata": {}, "source": [ "# Bring your own CNN: a trained U-Net from checkpoint to Earth Engine\n", "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)\n", "\n", - "Takes a U-Net trained with\n", + "Runs a U-Net you trained with\n", "[segmentation-models-pytorch](https://github.com/qubvel-org/segmentation_models.pytorch) (SMP)\n", - "and runs it inside Earth Engine on your imagery:\n", - "\n", - "1. check the model and find its **halo** (the border each tile needs);\n", - "2. plan the tiles;\n", - "3. upload the weights once, as Earth Engine table assets;\n", - "4. predict one test tile and compare it with PyTorch;\n", - "5. export a region.\n", - "\n", - "- **Offline first.** As shipped, it uses a small demo U-Net with random weights and only builds\n", - " the Earth Engine computations; nothing is sent.\n", - "- **Then live, step by step.** Set your checkpoint, project and input, set `RUN_LIVE = True`,\n", - " and turn on each Earth Engine step with its own switch. Each step uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate.\n", - "- **What is checked.** The building blocks agree with PyTorch at their validated settings, and\n", - " one trained ResNet-18 U-Net, prepared this way, matched PyTorch to 1.2e-5 in Earth Engine on a\n", - " 64 × 64 input. This tiled recipe has not been run live as a whole: step 9 checks it for your\n", - " model. Agreement with PyTorch is not accuracy." + "inside Earth Engine, on your own imagery.\n", + "\n", + "**How to use it**\n", + "\n", + "1. **Try it as it is:** **Runtime → Run all**. It uses a small demo U-Net with random weights\n", + " and sends nothing to Earth Engine. Each cell prints what it did.\n", + "2. **Set up your model:** in step 2, fill in your checkpoint, project, input image and region,\n", + " and set `USE_EARTH_ENGINE = True`. **Run all** again: it signs you in and checks your model,\n", + " and stops at step 8 until your weights are uploaded. Nothing is created yet.\n", + "3. **Run the Earth Engine steps one at a time, in this order.** Each has a checkbox in its own\n", + " cell: tick it, run that cell, wait until it prints \"Done\", untick it, and run the cells\n", + " below it (or **Run all** again).\n", + " - step 7: upload the weights (once);\n", + " - step 9: export one test tile;\n", + " - step 10: compare that tile with PyTorch;\n", + " - step 11: export the whole region.\n", + "\n", + "**What it creates in your project:** the weights as table assets (step 5 says how many), and\n", + "one image asset each for the test tile and the region.\n", + "\n", + "**Cost:** steps 7 to 11 use your project's Earth Engine quota (EECU time). That is free on a\n", + "noncommercial project; a commercial project pays its plan's rate. Their cost has not been\n", + "measured for this notebook.\n", + "\n", + "**What has been checked:** TensorEE's building blocks (its versions of each layer) agree with\n", + "PyTorch at their validated settings, and one trained ResNet-18 U-Net, prepared this way,\n", + "matched PyTorch to 1.2e-5 in Earth Engine on a 64 × 64 input. This tiled recipe has not been\n", + "run live as a whole: step 10 checks it for your model. Agreement with PyTorch is not accuracy." ] }, { "cell_type": "code", - "id": "b29a1f17", + "execution_count": null, + "id": "c4686e47", "metadata": {}, + "outputs": [], "source": [ "# @title 1. Install\n", - "# Install from the repository: support for more layers is added there between releases.\n", - "%pip install --quiet \"tensor-ee[segwater] @ git+https://github.com/noelEOS/tensor-ee\"" - ], - "execution_count": null, - "outputs": [] + "# Installs from the repository, which gets support for more layers between releases.\n", + "%pip install --quiet \"tensor-ee[segwater] @ git+https://github.com/noelEOS/tensor-ee\"\n", + "print(\"TensorEE is installed.\")" + ] }, { "cell_type": "markdown", - "id": "b0fec902", + "id": "ca2d2707", "metadata": {}, "source": [ "## 2. Settings\n", "\n", - "- **Model:** `CHECKPOINT` is your `state_dict` file (in Colab, upload it in the **Files** panel).\n", - " `None` = the offline demo. `ENCODER`, `BANDS` and `CLASSES` must match your training.\n", - "- **Input:** `INPUT_IMAGE` is an Earth Engine image with bands `BANDS`, already scaled as in\n", - " training. `REGION_X`, `REGION_Y` are the region's upper-left corner in `CRS` metres; the image\n", - " must also cover one halo around it (step 5 prints the width).\n", - "- **Earth Engine steps** run only when `RUN_LIVE` and their own switch are both `True`. You can\n", - " change `RUN_LIVE` at any time and run from the top again." + "- Each line says whether you need to change it. As shipped, the notebook runs offline with the\n", + " demo model, and none of these need changing.\n", + "- Your input image must already be in Earth Engine, with your model's bands, scaled as in\n", + " training. It must cover the region plus a border on every side (step 6 prints the exact area)." ] }, { "cell_type": "code", - "id": "2f178b02", + "execution_count": null, + "id": "a21e978b", "metadata": {}, + "outputs": [], "source": [ - "PROJECT = \"YOUR_EE_PROJECT\" # your Cloud project registered for Earth Engine\n", - "RUN_LIVE = False # True: sign in; the switches below choose the Earth Engine steps\n", - "\n", - "# Your model\n", - "CHECKPOINT = None # for example \"my_unet.pth\"; None = demo model, offline only\n", - "ENCODER = \"resnet34\"\n", - "BANDS = (\"VV\", \"VH\") # input bands, in training order\n", - "CLASSES = 2\n", - "\n", - "# Your input and region\n", - "INPUT_IMAGE = \"projects/YOUR_EE_PROJECT/assets/YOUR_PREPARED_INPUT\"\n", - "CRS = \"EPSG:32651\"\n", - "REGION_X, REGION_Y = 286030, 2640860 # upper-left corner of the region, CRS metres\n", - "PIXEL = 10 # metres per pixel\n", - "TILE = 1024 # region tile side, pixels\n", - "REGION_COLUMNS, REGION_ROWS = 2, 1 # region size, in tiles\n", - "TEST_TILE = 256 # side of the test tile (step 8)\n", - "\n", - "# What the live steps create\n", - "WEIGHTS_MANIFEST = \"my_unet_weights.json\" # local record of the uploaded weights: keep it\n", - "TEST_ASSET = \"my_unet_test_tile\"\n", - "REGION_ASSET = \"my_unet_region\"\n", - "\n", - "# Earth Engine steps\n", - "UPLOAD_WEIGHTS = False # step 7\n", - "PREDICT_TEST_TILE = False # step 8\n", - "COMPARE_TEST_TILE = False # step 9\n", - "EXPORT_REGION = False # step 10" - ], - "execution_count": null, - "outputs": [] + "# Earth Engine\n", + "PROJECT = \"YOUR_EE_PROJECT\" # change to your Cloud project ID, as in the Code Editor\n", + "USE_EARTH_ENGINE = False # leave False to try the demo offline; True once your model is set below\n", + "\n", + "# Your model: must match how it was trained\n", + "CHECKPOINT = None # change to your saved state_dict, for example \"my_unet.pth\"\n", + " # (in Colab, upload it in the Files panel on the left)\n", + "ENCODER = \"resnet34\" # change to the encoder_name you trained with\n", + "BANDS = (\"VV\", \"VH\") # change to your input bands, in training order\n", + "CLASSES = 2 # change to your number of classes\n", + "\n", + "# Your input image and region (the values here are an example in Taiwan)\n", + "INPUT_IMAGE = \"projects/YOUR_EE_PROJECT/assets/YOUR_PREPARED_INPUT\" # change to your image\n", + "CRS = \"EPSG:32651\" # change to a projection in metres, such as your UTM zone\n", + "REGION_X, REGION_Y = 286030, 2640860 # change to the region's upper-left corner, in CRS metres\n", + "REGION_COLUMNS, REGION_ROWS = 2, 1 # change to the region's size, in tiles of TILE pixels\n", + "\n", + "# Sizes: leave as they are unless your pixels are not 10 m or an export runs out of memory\n", + "PIXEL = 10 # metres per pixel\n", + "TILE = 1024 # side of each region tile, in pixels; larger tiles waste less on borders\n", + "TEST_TILE = 256 # side of the test tile, in pixels\n", + "\n", + "# Names for what the Earth Engine steps create: leave as they are, or rename\n", + "WEIGHTS_RECORD = \"my_unet_weights.json\" # local file listing your uploaded weight tables: keep it\n", + "TEST_ASSET = \"my_unet_test_tile\" # image asset for the test tile; an existing one is not overwritten\n", + "REGION_ASSET = \"my_unet_region\" # image asset for the region\n", + "\n", + "if CHECKPOINT is None:\n", + " print(\"Settings saved. No CHECKPOINT yet, so the notebook uses the demo model, offline.\")\n", + "elif not USE_EARTH_ENGINE:\n", + " print(f\"Settings saved. Model: {CHECKPOINT}. Offline until you set USE_EARTH_ENGINE = True.\")\n", + "else:\n", + " print(f\"Settings saved. Model: {CHECKPOINT}. Earth Engine project: {PROJECT}.\")" + ] }, { "cell_type": "code", - "id": "6193d146", - "metadata": {}, + "execution_count": null, + "id": "717f033c", + "metadata": { + "cellView": "form" + }, + "outputs": [], "source": [ + "# @title 3. Connect to Earth Engine, or stay offline\n", + "# @markdown Offline, the notebook prepares every Earth Engine request but sends none. You can\n", + "# @markdown change `USE_EARTH_ENGINE` at any time and run from the top again.\n", + "import ee\n", + "from ee.apitestcase import ApiTestCase\n", + "\n", + "\n", "class Stop(Exception):\n", " \"\"\"Stops the notebook with a one-line message instead of a traceback.\"\"\"\n", "\n", @@ -106,70 +132,66 @@ " return [f\"Stopped: {self}\"]\n", "\n", "\n", - "def live_step(what, switch=True):\n", - " \"\"\"True when an Earth Engine step should run.\"\"\"\n", - " if not (RUN_LIVE and switch):\n", - " print(f\"Skipped: {what} ({'RUN_LIVE' if not RUN_LIVE else 'its switch'} is False).\")\n", + "def live_step(checkbox=None, ticked=True):\n", + " \"\"\"True when an Earth Engine step should run; otherwise says why it does not.\"\"\"\n", + " if not USE_EARTH_ENGINE:\n", + " print(\"Not run: the notebook is offline (USE_EARTH_ENGINE is False in step 2).\")\n", + " return False\n", + " if not ticked:\n", + " print(f\"Not run: tick {checkbox} above when you are ready, then run this cell.\")\n", " return False\n", " if not PROJECT or PROJECT == \"YOUR_EE_PROJECT\":\n", - " raise Stop(\"set PROJECT to your Cloud project ID in the settings cell.\")\n", + " raise Stop(\"set PROJECT to your Cloud project ID in step 2.\")\n", " if CHECKPOINT is None:\n", - " raise Stop(\"live steps need your trained checkpoint: set CHECKPOINT.\")\n", - " return True" - ], - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "2cfc7c5f", - "metadata": {}, - "source": [ - "# @title 3. Connect: offline signatures, or your project\n", - "# @markdown Offline, Earth Engine's own test helper lets the notebook build computations without\n", - "# @markdown credentials. Switching to live undoes it, so no restart is needed.\n", - "import ee\n", - "from ee.apitestcase import ApiTestCase\n", + " raise Stop(\"the Earth Engine steps need your trained model: set CHECKPOINT in step 2.\")\n", + " return True\n", + "\n", "\n", - "if \"offline\" in globals(): # an earlier offline run in this session: undo it\n", + "# Offline, Earth Engine's own test helper stands in for a signed-in session.\n", + "if \"offline\" in globals(): # undo an earlier offline run in this session\n", " offline.tearDown()\n", " ee.data._install_cloud_api_resource = _install_cloud_api_resource\n", " ee.Reset()\n", " del offline\n", - "if not RUN_LIVE:\n", + "if USE_EARTH_ENGINE:\n", + " print(\"Earth Engine mode: the next cell signs you in.\")\n", + "else:\n", " _install_cloud_api_resource = ee.data._install_cloud_api_resource\n", " offline = ApiTestCase()\n", " offline.setUp()\n", - " print(\"Offline: computations are built, nothing is sent.\")" - ], - "execution_count": null, - "outputs": [] + " print(\"Offline mode: nothing will be sent to Earth Engine or created in a project.\")" + ] }, { "cell_type": "code", - "id": "695f4df2", + "execution_count": null, + "id": "70296347", "metadata": { "tags": [ "live" ] }, + "outputs": [], "source": [ - "if live_step(\"sign in\"):\n", + "# @title Sign in (Earth Engine mode only)\n", + "# @markdown The first time, follow the link with the Google account you use for Earth Engine.\n", + "if live_step():\n", " ee.Authenticate()\n", " ee.Initialize(project=PROJECT)\n", - " print(\"Signed in to\", PROJECT)" - ], - "execution_count": null, - "outputs": [] + " print(f\"Earth Engine is ready — signed in to {PROJECT}.\")" + ] }, { "cell_type": "code", - "id": "84bcd7d7", + "execution_count": null, + "id": "6efd6d01", "metadata": {}, + "outputs": [], "source": [ "# @title 4. Load the model\n", - "# @markdown `None`: a shallow demo U-Net with random weights. A checkpoint is loaded strictly:\n", - "# @markdown a missing or extra weight name means the settings do not match your training.\n", + "# @markdown Without a `CHECKPOINT`, a small demo U-Net with random weights. Your checkpoint must\n", + "# @markdown match `ENCODER`, `BANDS` and `CLASSES`: if not, this cell fails with an error listing\n", + "# @markdown the weight names that do not fit.\n", "import torch\n", "import segmentation_models_pytorch as smp\n", "\n", @@ -187,66 +209,78 @@ " torch.manual_seed(20261020)\n", " model = make_model(encoder_depth=3, decoder_channels=(16, 8, 4))\n", " tile, test_tile = 64, 64 # small tiles for the demo\n", - " print(\"Demo model: a shallow U-Net with random weights.\")\n", + " print(\"Loaded the demo model: a small U-Net with random weights.\")\n", "else:\n", " model = make_model()\n", " model.load_state_dict(torch.load(CHECKPOINT, map_location=\"cpu\", weights_only=True), strict=True)\n", " tile, test_tile = TILE, TEST_TILE\n", - " print(\"Loaded\", CHECKPOINT)" - ], - "execution_count": null, - "outputs": [] + " print(f\"Loaded your model from {CHECKPOINT}.\")" + ] }, { "cell_type": "code", - "id": "d4b2bff3", + "execution_count": null, + "id": "4e6dbb6c", "metadata": {}, + "outputs": [], "source": [ - "# @title 5. Check support and find the halo\n", - "# @markdown The halo is how far each output pixel sees into the input. Every tile reads the\n", - "# @markdown tile plus a halo on each side, then drops it: for a ResNet-34 U-Net, 544 pixels,\n", - "# @markdown so a 1,024-pixel tile reads 2,112 × 2,112. Larger tiles waste less.\n", - "from tensor_ee import check_compatibility\n", + "# @title 5. Check the model, its border and its weights\n", + "# @markdown - Checks that TensorEE has a building block for every layer of the model.\n", + "# @markdown - Each tile is computed from the tile plus an extra border of input on every side\n", + "# @markdown (the halo), so that its edges come out as they would on the whole image; the\n", + "# @markdown border is then dropped. For a ResNet-34 U-Net it is 544 pixels, so a 1,024-pixel\n", + "# @markdown tile reads 2,112 × 2,112.\n", + "# @markdown - Counts the weights step 7 uploads. BatchNorm is first merged into the convolutions.\n", + "import math\n", + "\n", + "from tensor_ee import check_compatibility, effective_convolutions\n", "from tensor_ee.receptive_field import analyze_receptive_field\n", "from tensor_ee.smp import adapt_smp_model\n", "\n", "\n", "def prepare(height, width):\n", - " \"\"\"The model adapted to one input size; stops unless every operation is supported.\"\"\"\n", + " \"\"\"A copy of the model for one input size; stops unless TensorEE can run every layer.\"\"\"\n", " shape = (1, len(BANDS), height, width)\n", " adapted = adapt_smp_model(model, input_shape=shape)\n", " report = check_compatibility(adapted, shape)\n", " if report.status != \"supported\":\n", - " raise Stop(f\"TensorEE does not support this model: {report.to_dict()['blockers'][:3]} \"\n", - " \"(notebook 02 explains reports).\")\n", + " raise Stop(f\"TensorEE cannot run this model yet: {report.to_dict()['blockers'][:3]} \"\n", + " \"(notebook 02 explains why a model is refused).\")\n", " return adapted\n", "\n", "\n", "probe = prepare(64, 64)\n", "support = analyze_receptive_field(probe)\n", "halo, alignment = support.safe_halo, support.alignment\n", - "print(\"Supported; the prepared model equals yours on a test input:\",\n", - " probe.tensor_ee_smp_metadata[\"torch_equal\"])\n", - "print(f\"Halo {halo} px ({halo * PIXEL / 1000:g} km) per side; sizes in multiples of {alignment} px.\")" - ], - "execution_count": null, - "outputs": [] + "parameters = effective_convolutions(probe)\n", + "values = sum(p.weights.size + (0 if p.bias is None else p.bias.size) for p in parameters)\n", + "tables = math.ceil(values * 14 / 7_000_000)\n", + "\n", + "print(\"Supported: TensorEE has a building block for every layer of this model.\")\n", + "if probe.tensor_ee_smp_metadata[\"torch_equal\"]:\n", + " print(\"The copy prepared for Earth Engine gives exactly your model's output on a test input.\")\n", + "print(f\"Border (halo): {halo} px ({halo * PIXEL / 1000:g} km) on each side. \"\n", + " f\"Tile sides must be multiples of {alignment} px.\")\n", + "print(f\"Weights to upload: {values:,} numbers, about {tables} table asset(s).\")" + ] }, { "cell_type": "code", - "id": "aa4d0e83", + "execution_count": null, + "id": "79dd01f2", "metadata": {}, + "outputs": [], "source": [ "# @title 6. Plan the tiles\n", - "# @markdown One test tile (step 8) and the region (step 10), each with a full halo of real input.\n", + "# @markdown The test tile is the region's upper-left tile. Each tile reads its border of real input.\n", "from tensor_ee import Grid\n", "from tensor_ee.converter import PixelWindow, plan_tiles\n", "\n", "\n", "def plan_region(side, columns, rows):\n", - " for name, value in ((\"tile side\", side), (\"halo\", halo)):\n", + " for name, value in ((\"tile side\", side), (\"border\", halo)):\n", " if value % alignment:\n", - " raise Stop(f\"the {name} ({value}) must be a multiple of {alignment} pixels.\")\n", + " raise Stop(f\"the {name} ({value} px) must be a multiple of {alignment} px.\")\n", " padded = Grid(CRS, (PIXEL, 0, REGION_X - halo * PIXEL, 0, -PIXEL, REGION_Y + halo * PIXEL),\n", " columns * side + 2 * halo, rows * side + 2 * halo)\n", " region = PixelWindow(halo, halo, columns * side, rows * side)\n", @@ -254,85 +288,83 @@ "\n", "\n", "def describe(plan, name):\n", - " context = plan.tiles[0].context\n", + " context, side = plan.tiles[0].context, plan.core_size[0]\n", " ratio = sum(t.context.pixels for t in plan.tiles) / plan.roi.pixels\n", - " print(f\"{name}: {len(plan.tiles)} tile(s) of {plan.core_size[0]} px, each reading \"\n", - " f\"{context.width} x {context.height} ({ratio:.2f} input px per output px); input \"\n", - " f\"{plan.global_grid.width} x {plan.global_grid.height} px from \"\n", - " f\"({plan.global_grid.transform[2]}, {plan.global_grid.transform[5]}).\")\n", + " print(f\"{name}: {len(plan.tiles)} tile(s) of {side} × {side} px, each reading \"\n", + " f\"{context.width} × {context.height} px with its border \"\n", + " f\"({ratio:.1f} input pixels per output pixel).\")\n", "\n", "\n", "test_plan = plan_region(test_tile, 1, 1)\n", "region_plan = plan_region(tile, REGION_COLUMNS, REGION_ROWS)\n", "describe(test_plan, \"Test tile\")\n", - "describe(region_plan, \"Region\")" - ], - "execution_count": null, - "outputs": [] + "describe(region_plan, \"Region\")\n", + "needed = region_plan.global_grid\n", + "print(f\"Your input image must cover {needed.width} × {needed.height} px from the upper-left \"\n", + " f\"corner ({needed.transform[2]}, {needed.transform[5]}) in {CRS}.\")\n", + "if CHECKPOINT is None:\n", + " print(\"(The demo uses 64 px tiles, not TILE and TEST_TILE.)\")" + ] }, { "cell_type": "code", - "id": "42d821f2", - "metadata": {}, - "source": [ - "# @title 7. Upload the weights once\n", - "# @markdown - BatchNorm is folded into the convolutions, and the numbers go into table assets\n", - "# @markdown `projects/PROJECT/assets/tensor_ee_weights_`, one export task each (two at a time).\n", - "# @markdown - `WEIGHTS_MANIFEST` records them locally: keep it (in Colab, download it).\n", - "# @markdown - Cost not recorded; it uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate. The trained ResNet-18 U-Net (14.3 M numbers) took 29 tables.\n", - "import math\n", - "\n", - "from tensor_ee import effective_convolutions\n", - "\n", - "parameters = effective_convolutions(probe)\n", - "values = sum(p.weights.size + (0 if p.bias is None else p.bias.size) for p in parameters)\n", - "print(f\"{len(parameters)} folded convolutions, {values:,} numbers: \"\n", - " f\"about {math.ceil(values * 14 / 7_000_000)} tables.\")" - ], "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "73b9c3ce", + "id": "ad3edbbe", "metadata": { "tags": [ "live" ] }, + "outputs": [], "source": [ + "# @title 7. Upload the weights (once)\n", + "# @markdown - Creates the weight table assets, under `projects/PROJECT/assets/tensor_ee_weights_`,\n", + "# @markdown with one export task each, two at a time.\n", + "# @markdown - Writes `WEIGHTS_RECORD`, the local list of those tables, which the prediction reads.\n", + "# @markdown Keep a copy: in Colab, download it from the Files panel, since Colab's files are\n", + "# @markdown lost when its runtime is reset.\n", + "# @markdown - Cost not measured. For scale, a trained ResNet-18 U-Net (14.3 M numbers) took 29 tables.\n", + "UPLOAD_WEIGHTS = False # @param {type:\"boolean\"}\n", + "\n", "import json\n", "import uuid\n", - "from collections import Counter\n", "from pathlib import Path\n", "\n", "from tensor_ee import export_weights\n", "\n", - "if live_step(\"upload the weights\", UPLOAD_WEIGHTS):\n", - " manifest_path = Path(WEIGHTS_MANIFEST)\n", - " if manifest_path.exists():\n", - " raise Stop(f\"{manifest_path} exists, so the weights are already uploaded.\")\n", + "if live_step(\"UPLOAD_WEIGHTS\", UPLOAD_WEIGHTS):\n", + " record_path = Path(WEIGHTS_RECORD)\n", + " if record_path.exists():\n", + " raise Stop(f\"the weights are already uploaded ({record_path} exists): untick \"\n", + " \"UPLOAD_WEIGHTS and go on.\")\n", " namespace = f\"projects/{PROJECT}/assets/tensor_ee_weights_{uuid.uuid4().hex[:12]}\"\n", - "\n", - " def progress(manifest):\n", - " states = Counter(shard[\"status\"] for shard in manifest[\"shards\"])\n", - " print(json.dumps({\"status\": manifest[\"status\"], \"tables\": dict(states)}), flush=True)\n", - "\n", - " export_weights(parameters, namespace, manifest_path, on_update=progress)\n", - " print(\"Uploaded; manifest written to\", manifest_path)" - ], - "execution_count": null, - "outputs": [] + " reported = []\n", + "\n", + " def progress(record):\n", + " done = sum(shard[\"status\"] == \"COMPLETED\" for shard in record[\"shards\"])\n", + " if not reported or done != reported[-1]:\n", + " reported.append(done)\n", + " print(f\" {done} of about {tables} table(s) uploaded...\", flush=True)\n", + "\n", + " print(\"Uploading; follow it here or in the Tasks tab: https://code.earthengine.google.com/tasks\")\n", + " export_weights(parameters, namespace, record_path, on_update=progress)\n", + " print(f\"Done: the weights are uploaded, and {record_path} lists them. Untick \"\n", + " \"UPLOAD_WEIGHTS, then run on from step 8.\")" + ] }, { "cell_type": "code", - "id": "3c1fd62e", + "execution_count": null, + "id": "efe612f2", "metadata": {}, + "outputs": [], "source": [ - "# @title 8. Predict the test tile\n", - "# @markdown Builds the computation (offline, with stand-in weight tables). Output bands: `score_k`\n", - "# @markdown (logits), `p_k` (probabilities), `class_index` and, for two classes, `binary_class`\n", - "# @markdown (`p_1` ≥ 0.5). Pixels missing any input band get no prediction.\n", + "# @title 8. Set up the prediction (nothing runs yet)\n", + "# @markdown - Connects the model to your input image and uploaded weights, for the test tile.\n", + "# @markdown Offline, placeholder weights stand in: the prediction is complete but cannot run.\n", + "# @markdown - Output bands: `score_k` (the model's raw scores, or logits), `p_k` (probabilities),\n", + "# @markdown `class_index` and, for two classes, `binary_class` (1 where `p_1` ≥ 0.5).\n", + "# @markdown - Pixels missing any input band get no prediction.\n", "import json\n", "from pathlib import Path\n", "\n", @@ -344,11 +376,13 @@ "PROBABILITIES = tuple(f\"p_{k}\" for k in range(CLASSES))\n", "\n", "\n", - "def weights_manifest():\n", - " if RUN_LIVE and Path(WEIGHTS_MANIFEST).exists():\n", - " return json.loads(Path(WEIGHTS_MANIFEST).read_text())\n", - " if RUN_LIVE:\n", - " raise Stop(f\"{WEIGHTS_MANIFEST} not found: upload the weights first (step 7).\")\n", + "def weights_record():\n", + " \"\"\"The uploaded weight tables, or offline placeholders with the same layout.\"\"\"\n", + " if USE_EARTH_ENGINE and Path(WEIGHTS_RECORD).exists():\n", + " return json.loads(Path(WEIGHTS_RECORD).read_text())\n", + " if USE_EARTH_ENGINE:\n", + " raise Stop(f\"{WEIGHTS_RECORD} not found. Upload the weights first (tick UPLOAD_WEIGHTS \"\n", + " \"in step 7), or, if you uploaded them before, put that file back in this folder.\")\n", " return {\"format\": \"native_ee_featurecollection_weights_v2\", \"status\": \"ready\",\n", " \"layers\": {p.key: {\"shape\": list(p.shape), \"dtype\": p.dtype, \"bias\": p.bias is not None,\n", " \"source\": p.source, \"sha256\": p.sha256,\n", @@ -357,7 +391,7 @@ "\n", "\n", "def input_image():\n", - " if RUN_LIVE:\n", + " if USE_EARTH_ENGINE:\n", " return ee.Image(INPUT_IMAGE)\n", " return ee.Image.constant([0] * len(BANDS)).rename(list(BANDS)) # offline stand-in\n", "\n", @@ -366,8 +400,8 @@ " context = plan.tiles[0].context\n", " adapted = prepare(context.height, context.width)\n", " if analyze_receptive_field(adapted).safe_halo != halo:\n", - " raise Stop(\"the halo changed with the input size; revise the plan.\")\n", - " provider = MatchedWeightProvider(adapted, weights_manifest())\n", + " raise Stop(\"the border changed with the input size; please report it as an issue.\")\n", + " provider = MatchedWeightProvider(adapted, weights_record())\n", " source = input_image().select(list(BANDS)).toFloat().reproject(plan.global_grid.projection())\n", " valid = Raster(source.mask().reduce(ee.Reducer.min()).rename(\"valid\"), (\"valid\",), plan.global_grid)\n", " logits = compile_tiled_model(adapted, source, BANDS, plan, weight_provider=provider)\n", @@ -390,60 +424,80 @@ "\n", "with deep_recursion():\n", " test_prediction = build_prediction(test_plan)\n", - " test_graph_bytes = len(test_prediction.serialize().encode())\n", - "print(f\"Test tile computation: {test_graph_bytes:,} bytes of Earth Engine expression.\")" - ], - "execution_count": null, - "outputs": [] + " test_graph_bytes = len(test_prediction.serialize().encode()) # checks it can be sent\n", + "if USE_EARTH_ENGINE:\n", + " print(\"The test-tile prediction is set up with your uploaded weights. Step 9 runs it.\")\n", + "else:\n", + " print(\"The test-tile prediction is set up offline, with placeholder weights: complete, \"\n", + " \"but it cannot run.\")" + ] }, { "cell_type": "code", - "id": "41ea3bf8", + "execution_count": null, + "id": "78c0ce07", "metadata": { "tags": [ "live" ] }, + "outputs": [], "source": [ - "# @title Export the test tile\n", - "# @markdown One batch export, one new image asset (the tile, not its halo). Cost not measured;\n", - "# @markdown it uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate. For scale, notebook 01's 224 × 224 chip used about 65 EECU-seconds.\n", - "# @markdown If it runs out of memory or time, lower `TEST_TILE` rather than resubmitting.\n", + "# @title 9. Export the test tile\n", + "# @markdown - Computes the test tile in a batch export and saves it as the image asset `TEST_ASSET`\n", + "# @markdown (the tile only, not its border).\n", + "# @markdown - Cost not measured. For scale, notebook 01's 224 × 224 chip used about 65 EECU-seconds.\n", + "# @markdown - If it fails for lack of memory or time, lower `TEST_TILE` in step 2 and run from the top.\n", + "EXPORT_TEST_TILE = False # @param {type:\"boolean\"}\n", + "\n", "import time\n", "\n", + "TASKS_URL = \"https://code.earthengine.google.com/tasks\"\n", + "\n", "\n", "def wait(task, minutes=60):\n", - " deadline = time.time() + minutes * 60\n", + " \"\"\"Wait for an export, with a line every 5 minutes; stop if it fails.\"\"\"\n", + " started, polls = time.time(), 0\n", " while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", - " if time.time() > deadline:\n", - " raise Stop(\"still running; follow it in the Tasks tab, or run this cell again later.\")\n", + " if time.time() - started > minutes * 60:\n", + " raise Stop(f\"still running after {minutes} minutes. It carries on in Earth Engine: \"\n", + " f\"follow it at {TASKS_URL}, and go on once it has finished.\")\n", " time.sleep(30)\n", + " polls += 1\n", + " if polls % 10 == 0:\n", + " print(f\" still running, {polls // 2} min...\", flush=True)\n", " if state != \"COMPLETED\":\n", - " raise Stop(f\"the export ended {state}: {task.status().get('error_message', '')}\")\n", - " print(\"Task\", task.id, \"completed\")\n", + " raise Stop(f\"the export ended {state}: {task.status().get('error_message', '')} \"\n", + " f\"(see {TASKS_URL})\")\n", "\n", "\n", - "if live_step(\"export the test tile\", PREDICT_TEST_TILE):\n", + "if live_step(\"EXPORT_TEST_TILE\", EXPORT_TEST_TILE):\n", " test_task = export_task(test_prediction, test_plan, TEST_ASSET, \"tensor_ee_test_tile\")\n", " test_task.start()\n", - " wait(test_task)" - ], - "execution_count": null, - "outputs": [] + " print(f\"Export started; follow it here or at {TASKS_URL}\", flush=True)\n", + " wait(test_task)\n", + " print(f\"Done: the test tile is saved as projects/{PROJECT}/assets/{TEST_ASSET}. Untick \"\n", + " \"EXPORT_TEST_TILE, then go on to step 10.\")" + ] }, { "cell_type": "code", - "id": "1a50cc0b", + "execution_count": null, + "id": "4b544941", "metadata": { "tags": [ "live" ] }, + "outputs": [], "source": [ - "# @title 9. Compare the test tile with PyTorch\n", - "# @markdown Downloads the tile's input (with halo) and the exported logits (2 small requests,\n", - "# @markdown no assets), runs your model locally, and reports the largest absolute difference\n", - "# @markdown (target 1e-4). The input must be complete; it is refused above 40 MB.\n", + "# @title 10. Compare the test tile with PyTorch\n", + "# @markdown - Run it once step 9 is done. It downloads the test tile's input (with its border)\n", + "# @markdown and the exported scores: 2 small requests, nothing created.\n", + "# @markdown - It then runs your model here on the same input and prints the largest difference\n", + "# @markdown (target: at most 1e-4). The input must have no gaps; above 40 MB it is refused.\n", + "COMPARE_WITH_PYTORCH = False # @param {type:\"boolean\"}\n", + "\n", "import numpy as np\n", "\n", "\n", @@ -456,66 +510,73 @@ " return np.stack([values[band] for band in bands]).astype(np.float32)\n", "\n", "\n", - "if live_step(\"compare the test tile with PyTorch\", COMPARE_TEST_TILE):\n", + "if live_step(\"COMPARE_WITH_PYTORCH\", COMPARE_WITH_PYTORCH):\n", " padded = test_plan.global_grid\n", " size = padded.width * padded.height * (len(BANDS) + 1) * 4\n", " if size > 40_000_000:\n", - " raise Stop(f\"the input is {size / 1e6:.0f} MB; lower TEST_TILE.\")\n", + " raise Stop(f\"the input is {size / 1e6:.0f} MB; lower TEST_TILE in step 2.\")\n", " source = ee.Image(INPUT_IMAGE).select(list(BANDS)).toFloat().reproject(padded.projection())\n", " valid = source.mask().reduce(ee.Reducer.min()).rename(\"valid\")\n", " pixels = download(source.unmask(0).addBands(valid), padded, (*BANDS, \"valid\"))\n", " if pixels[-1].min() < 1:\n", - " raise Stop(\"the input has gaps in the tile or its halo; move the region or fill them.\")\n", + " raise Stop(\"the input has gaps in the test tile or its border; move the region or fill them.\")\n", " exported = download(ee.Image(f\"projects/{PROJECT}/assets/{TEST_ASSET}\"), test_plan.roi_grid, SCORES)\n", "\n", " with torch.no_grad():\n", " reference = model(torch.from_numpy(pixels[None, :-1].copy()))[0].numpy()\n", " core = reference[:, halo:halo + test_tile, halo:halo + test_tile]\n", " error = np.abs(exported - core)\n", - " print(f\"Largest absolute logit difference: {error.max():.3e} (target 1e-4); \"\n", - " f\"mean {error.mean():.3e}; class disagreements \"\n", - " f\"{int((exported.argmax(0) != core.argmax(0)).sum())} of {core[0].size} pixels.\")\n", - " print(\"Within the target.\" if error.max() <= 1e-4 else\n", - " \"Above the target: please report it, with these numbers, as an issue.\")" - ], - "execution_count": null, - "outputs": [] + " print(f\"Largest difference in the scores: {error.max():.3e} (target 1e-4); \"\n", + " f\"mean {error.mean():.3e}; pixels with a different class: \"\n", + " f\"{int((exported.argmax(0) != core.argmax(0)).sum())} of {core[0].size}.\")\n", + " print(\"Done: within the target. Earth Engine agrees with PyTorch on this tile.\" if error.max() <= 1e-4\n", + " else \"Done, but above the target: please report it as an issue, with these numbers.\")" + ] }, { "cell_type": "code", - "id": "47fb0596", + "execution_count": null, + "id": "020c43c6", "metadata": { "tags": [ "live" ] }, + "outputs": [], "source": [ - "# @title 10. Export the region\n", - "# @markdown One batch export, one new image asset over the region (not its halo). Cost and\n", - "# @markdown capacity not measured at these tile sizes; it uses your project's Earth Engine quota (EECU time): free on a noncommercial project; a commercial project pays its plan's rate. Run the test tile first.\n", - "if live_step(\"export the region\", EXPORT_REGION):\n", + "# @title 11. Export the region\n", + "# @markdown - Run it once step 10 is within the target. One batch export, saved as the image\n", + "# @markdown asset `REGION_ASSET` (the region only, not its border).\n", + "# @markdown - Cost, memory use and time not measured at these tile sizes.\n", + "EXPORT_REGION = False # @param {type:\"boolean\"}\n", + "\n", + "if live_step(\"EXPORT_REGION\", EXPORT_REGION):\n", " with deep_recursion():\n", " region_task = export_task(build_prediction(region_plan), region_plan, REGION_ASSET,\n", " \"tensor_ee_region\")\n", " region_task.start()\n", - " wait(region_task, minutes=240)" - ], - "execution_count": null, - "outputs": [] + " print(f\"Export started; follow it here or at {TASKS_URL}\", flush=True)\n", + " wait(region_task, minutes=240)\n", + " print(f\"Done: the region is saved as projects/{PROJECT}/assets/{REGION_ASSET}.\")" + ] }, { "cell_type": "markdown", - "id": "d6897497", + "id": "3dee0517", "metadata": {}, "source": [ "## Good to know\n", "\n", - "- **Evidence:** step 9 is live agreement for your model on one tile; keep its numbers.\n", - " Accuracy is separate: evaluate the exports against your own labelled data.\n", + "- **Coming back later:** put your `WEIGHTS_RECORD` file back next to the notebook (in Colab,\n", + " upload it in the Files panel, like the checkpoint), set\n", + " `USE_EARTH_ENGINE = True`, and **Run all** with every checkbox unticked; then tick only the\n", + " step you need.\n", + "- **Evidence:** step 10 is agreement in Earth Engine for your model on one tile; keep its\n", + " numbers. Accuracy is separate: evaluate the exports against your own labelled data.\n", "- **Details:** [the guide this notebook follows](https://github.com/noelEOS/tensor-ee/blob/main/docs/own-model.md)\n", - " covers other encoders, irregular regions and validity masks.\n", - "- **Clean up:** weight tables and images stay in your project until you delete them (Code\n", - " Editor, **Assets** tab); keep the weight tables while you predict with them." + " covers other encoders, irregular regions and masks of valid pixels.\n", + "- **Clean up:** the weight tables and images stay in your project until you delete them (Code\n", + " Editor, **Assets** tab). Keep the weight tables while you predict with them." ] } ], diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py index 1bd17dc..caacbee 100644 --- a/tests/test_notebooks.py +++ b/tests/test_notebooks.py @@ -8,8 +8,9 @@ - The quickstart (01) has no offline mode of its own: the test installs the client's offline signatures, and after its settings cell supplies a project name and the place and scene that its live cells would find (the preset Sun Moon Lake chip). -- 03 runs with ``RUN_LIVE = False`` and its own offline signatures; switching it to live must - undo them without a restart. +- 03 runs with ``USE_EARTH_ENGINE = False`` and its own offline signatures; switching it to + Earth Engine must undo them without a restart. Each of its Earth Engine steps has its own + checkbox, off as shipped. - 02 is offline throughout. Notebooks that need segmentation-models-pytorch or timm skip by name without them. @@ -144,11 +145,14 @@ def status(self): def test_live_steps_in_03_check_the_switches_and_the_project(): cells = _code("03_bring_your_own_cnn.ipynb") - settings = [cell for cell in cells if re.search(r"^RUN_LIVE = False\b", _source(cell), re.M)] + settings = [cell for cell in cells if re.search(r"^USE_EARTH_ENGINE = False\b", _source(cell), re.M)] assert len(settings) == 1 and f'PROJECT = "{PLACEHOLDER_PROJECT}"' in _source(settings[0]) + checkboxes = [] for cell in cells: if _live(cell): assert "live_step(" in _source(cell) + checkboxes += re.findall(r'^([A-Z_]+) = False # @param \{type:"boolean"\}$', _source(cell), re.M) + assert checkboxes == ["UPLOAD_WEIGHTS", "EXPORT_TEST_TILE", "COMPARE_WITH_PYTORCH", "EXPORT_REGION"] @pytest.mark.parametrize("name", sorted(EXPECTED)) @@ -226,8 +230,8 @@ def _run(name, monkeypatch, *, after_settings=None, signatures=False, then=None) raise if after_settings and "@param" in _source(cell): namespace.update(after_settings) - if "RUN_LIVE" in namespace: - assert namespace["RUN_LIVE"] is False + if "USE_EARTH_ENGINE" in namespace: + assert namespace["USE_EARTH_ENGINE"] is False if then: then(namespace) finally: @@ -268,17 +272,17 @@ def test_bring_your_own_cnn_runs_offline_and_goes_live_without_a_restart(monkeyp import ee def switch_to_live(namespace): - # Rerunning the connect cell with RUN_LIVE = True undoes the offline signatures. + # Rerunning the connect cell with USE_EARTH_ENGINE = True undoes the offline signatures. connect = next((index, cell) for index, cell in enumerate(_cells("03_bring_your_own_cnn.ipynb")) if "ApiTestCase()" in _source(cell)) - namespace["RUN_LIVE"] = True + namespace["USE_EARTH_ENGINE"] = True offline = namespace["offline"] _exec("03_bring_your_own_cnn.ipynb", *connect, namespace) assert "offline" not in namespace assert ee.data._install_cloud_api_resource is namespace["_install_cloud_api_resource"] assert ee.data.getAlgorithms is offline.old_get_algorithms assert ee.data.getMapId is offline.old_get_map_id - namespace["RUN_LIVE"] = False + namespace["USE_EARTH_ENGINE"] = False namespace = _run("03_bring_your_own_cnn.ipynb", monkeypatch, then=switch_to_live) assert namespace["CHECKPOINT"] is None From 4e177c2ea6c5c93ca35fca4c59204b0fddaa710d Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:36:58 +0800 Subject: [PATCH 8/9] Notebook 03: one setting for how far to go, a live demo on public Sentinel-1, and finished steps skipped - 03: RUN_UP_TO replaces USE_EARTH_ENGINE and the four checkboxes. Each Run all goes up to that step; an uploaded weights record, an existing asset or an export still running is reused, not an error. Without a checkpoint the Earth Engine steps run the small random-weight demo on a public Sentinel-1 scene at Sun Moon Lake, rescaled as (dB + 20) / 20, with demo-named assets. The region is exported only after a comparison within the target. - All three notebooks install with !pip, so no restart note is printed. - tests: the ladder, the skips (with stand-in tasks) and the absence of restart text. - start-here: the new 03 description. --- docs/start-here.md | 8 +- notebooks/01_quickstart.ipynb | 2 +- notebooks/02_is_my_model_supported.ipynb | 2 +- notebooks/03_bring_your_own_cnn.ipynb | 316 +++++++++++++---------- tests/test_notebooks.py | 95 +++++-- 5 files changed, 271 insertions(+), 152 deletions(-) diff --git a/docs/start-here.md b/docs/start-here.md index b649611..10456fd 100644 --- a/docs/start-here.md +++ b/docs/start-here.md @@ -34,10 +34,10 @@ For anyone with an Earth Engine account and some Python who wants to try TensorE refused and why. - [Notebook 03, bring your own CNN](../notebooks/03_bring_your_own_cnn.ipynb) ([open in Colab](https://colab.research.google.com/github/noelEOS/tensor-ee/blob/main/notebooks/03_bring_your_own_cnn.ipynb)): - a trained segmentation-models-pytorch U-Net from checkpoint to an exported region. Try it - offline with a demo model first; then each Earth Engine step (upload the weights, export a - test tile, compare it with PyTorch, export the region) has its own checkbox, run one at a - time. + a trained segmentation-models-pytorch U-Net from checkpoint to an exported region. One + setting says how far a run goes: offline first, then upload the weights, export a test + tile, compare it with PyTorch, and export the region. A small random-weight demo model can + try every step on public Sentinel-1 before you bring your own. For details, see [Getting started](getting-started.md) and the [guides](README.md). diff --git a/notebooks/01_quickstart.ipynb b/notebooks/01_quickstart.ipynb index 97a88ce..a9505b2 100644 --- a/notebooks/01_quickstart.ipynb +++ b/notebooks/01_quickstart.ipynb @@ -29,7 +29,7 @@ "# @title 1. Install (about a minute)\n", "# The model's Earth Engine computation is in the repository, not on PyPI, so install from it.\n", "!git clone --quiet --depth 1 https://github.com/noelEOS/tensor-ee.git\n", - "%pip install --quiet ./tensor-ee folium\n", + "!pip install --quiet ./tensor-ee folium\n", "\n", "import sys\n", "from pathlib import Path\n", diff --git a/notebooks/02_is_my_model_supported.ipynb b/notebooks/02_is_my_model_supported.ipynb index 696f4df..7b7d107 100644 --- a/notebooks/02_is_my_model_supported.ipynb +++ b/notebooks/02_is_my_model_supported.ipynb @@ -30,7 +30,7 @@ "# Installs from the repository, which gets support for more layers between releases. The\n", "# \"swin\" extra brings segmentation-models-pytorch 0.5.0 and timm 1.0.30, the versions TensorEE\n", "# was checked with.\n", - "%pip install --quiet \"tensor-ee[swin] @ git+https://github.com/noelEOS/tensor-ee\"\n", + "!pip install --quiet \"tensor-ee[swin] @ git+https://github.com/noelEOS/tensor-ee\"\n", "print(\"TensorEE is installed.\")" ] }, diff --git a/notebooks/03_bring_your_own_cnn.ipynb b/notebooks/03_bring_your_own_cnn.ipynb index d32a734..a95ef92 100644 --- a/notebooks/03_bring_your_own_cnn.ipynb +++ b/notebooks/03_bring_your_own_cnn.ipynb @@ -15,25 +15,30 @@ "\n", "**How to use it**\n", "\n", - "1. **Try it as it is:** **Runtime → Run all**. It uses a small demo U-Net with random weights\n", - " and sends nothing to Earth Engine. Each cell prints what it did.\n", - "2. **Set up your model:** in step 2, fill in your checkpoint, project, input image and region,\n", - " and set `USE_EARTH_ENGINE = True`. **Run all** again: it signs you in and checks your model,\n", - " and stops at step 8 until your weights are uploaded. Nothing is created yet.\n", - "3. **Run the Earth Engine steps one at a time, in this order.** Each has a checkbox in its own\n", - " cell: tick it, run that cell, wait until it prints \"Done\", untick it, and run the cells\n", - " below it (or **Run all** again).\n", - " - step 7: upload the weights (once);\n", - " - step 9: export one test tile;\n", - " - step 10: compare that tile with PyTorch;\n", - " - step 11: export the whole region.\n", - "\n", - "**What it creates in your project:** the weights as table assets (step 5 says how many), and\n", - "one image asset each for the test tile and the region.\n", - "\n", - "**Cost:** steps 7 to 11 use your project's Earth Engine quota (EECU time). That is free on a\n", - "noncommercial project; a commercial project pays its plan's rate. Their cost has not been\n", - "measured for this notebook.\n", + "- **One setting decides how far a run goes:** `RUN_UP_TO`, in step 2. **Runtime → Run all**\n", + " goes up to that step and skips the steps already done, so you can move it one step further\n", + " and run all again.\n", + "- **Start offline**, as shipped: a small demo U-Net with random weights. Nothing is sent to\n", + " Earth Engine. Each cell prints what it did.\n", + "- **Try the Earth Engine steps with the demo:** set `PROJECT`, then move `RUN_UP_TO` on. The\n", + " demo reads a public Sentinel-1 scene at Sun Moon Lake, Taiwan, with small 64-pixel tiles.\n", + " Its random weights make a meaningless map: the point is to see each step work, and Earth\n", + " Engine agree with PyTorch.\n", + "- **Then your model:** set `CHECKPOINT`, your model's settings, `INPUT_IMAGE` and the region,\n", + " and start again from `\"sign in\"`.\n", + "\n", + "| `RUN_UP_TO` | adds | creates in your project |\n", + "|---|---|---|\n", + "| `\"offline\"` | everything, without Earth Engine | nothing |\n", + "| `\"sign in\"` | signs in and sets up the prediction | nothing |\n", + "| `\"upload weights\"` | step 7: uploads the weights | table assets (3 for the demo) |\n", + "| `\"export test tile\"` | step 9: exports one test tile | one image asset |\n", + "| `\"compare with PyTorch\"` | step 10: compares that tile with PyTorch | nothing |\n", + "| `\"export region\"` | step 11: exports the whole region | one image asset |\n", + "\n", + "**Cost:** the Earth Engine steps use your project's Earth Engine quota (EECU time). That is\n", + "free on a noncommercial project; a commercial project pays its plan's rate. Their cost has not\n", + "been measured for this notebook.\n", "\n", "**What has been checked:** TensorEE's building blocks (its versions of each layer) agree with\n", "PyTorch at their validated settings, and one trained ResNet-18 U-Net, prepared this way,\n", @@ -50,7 +55,7 @@ "source": [ "# @title 1. Install\n", "# Installs from the repository, which gets support for more layers between releases.\n", - "%pip install --quiet \"tensor-ee[segwater] @ git+https://github.com/noelEOS/tensor-ee\"\n", + "!pip install --quiet \"tensor-ee[segwater] @ git+https://github.com/noelEOS/tensor-ee\"\n", "print(\"TensorEE is installed.\")" ] }, @@ -62,7 +67,7 @@ "## 2. Settings\n", "\n", "- Each line says whether you need to change it. As shipped, the notebook runs offline with the\n", - " demo model, and none of these need changing.\n", + " demo model, and none of them need changing.\n", "- Your input image must already be in Earth Engine, with your model's bands, scaled as in\n", " training. It must cover the region plus a border on every side (step 6 prints the exact area)." ] @@ -70,23 +75,24 @@ { "cell_type": "code", "execution_count": null, - "id": "a21e978b", + "id": "27953840", "metadata": {}, "outputs": [], "source": [ - "# Earth Engine\n", + "# How far this run goes. Each Run all goes up to this step and skips the steps already done.\n", + "RUN_UP_TO = \"offline\" # @param [\"offline\", \"sign in\", \"upload weights\", \"export test tile\", \"compare with PyTorch\", \"export region\"]\n", "PROJECT = \"YOUR_EE_PROJECT\" # change to your Cloud project ID, as in the Code Editor\n", - "USE_EARTH_ENGINE = False # leave False to try the demo offline; True once your model is set below\n", "\n", "# Your model: must match how it was trained\n", - "CHECKPOINT = None # change to your saved state_dict, for example \"my_unet.pth\"\n", - " # (in Colab, upload it in the Files panel on the left)\n", + "CHECKPOINT = None # leave None for the demo model; or your saved state_dict, for\n", + " # example \"my_unet.pth\" (in Colab, upload it in the Files panel)\n", "ENCODER = \"resnet34\" # change to the encoder_name you trained with\n", "BANDS = (\"VV\", \"VH\") # change to your input bands, in training order\n", "CLASSES = 2 # change to your number of classes\n", "\n", - "# Your input image and region (the values here are an example in Taiwan)\n", - "INPUT_IMAGE = \"projects/YOUR_EE_PROJECT/assets/YOUR_PREPARED_INPUT\" # change to your image\n", + "# Your input image and region (the defaults are the demo's: Sun Moon Lake, Taiwan)\n", + "INPUT_IMAGE = \"COPERNICUS/S1_GRD/S1A_IW_GRDH_1SDV_20240103T100134_20240103T100159_051941_06469E_B0BD\"\n", + " # change to your image; the default is a public Sentinel-1 scene\n", "CRS = \"EPSG:32651\" # change to a projection in metres, such as your UTM zone\n", "REGION_X, REGION_Y = 286030, 2640860 # change to the region's upper-left corner, in CRS metres\n", "REGION_COLUMNS, REGION_ROWS = 2, 1 # change to the region's size, in tiles of TILE pixels\n", @@ -96,17 +102,17 @@ "TILE = 1024 # side of each region tile, in pixels; larger tiles waste less on borders\n", "TEST_TILE = 256 # side of the test tile, in pixels\n", "\n", - "# Names for what the Earth Engine steps create: leave as they are, or rename\n", + "# Names for what the Earth Engine steps create (the demo uses its own, starting tensor_ee_demo_).\n", + "# An existing asset is reused, not overwritten: use new names when you change the model or region.\n", "WEIGHTS_RECORD = \"my_unet_weights.json\" # local file listing your uploaded weight tables: keep it\n", - "TEST_ASSET = \"my_unet_test_tile\" # image asset for the test tile; an existing one is not overwritten\n", + "TEST_ASSET = \"my_unet_test_tile\" # image asset for the test tile\n", "REGION_ASSET = \"my_unet_region\" # image asset for the region\n", "\n", - "if CHECKPOINT is None:\n", - " print(\"Settings saved. No CHECKPOINT yet, so the notebook uses the demo model, offline.\")\n", - "elif not USE_EARTH_ENGINE:\n", - " print(f\"Settings saved. Model: {CHECKPOINT}. Offline until you set USE_EARTH_ENGINE = True.\")\n", + "model_name = CHECKPOINT or \"the demo model (random weights)\"\n", + "if RUN_UP_TO == \"offline\":\n", + " print(f\"Settings saved. Model: {model_name}. Offline: nothing will be sent to Earth Engine.\")\n", "else:\n", - " print(f\"Settings saved. Model: {CHECKPOINT}. Earth Engine project: {PROJECT}.\")" + " print(f'Settings saved. Model: {model_name}. This run goes up to \"{RUN_UP_TO}\" in project {PROJECT}.')" ] }, { @@ -120,10 +126,13 @@ "source": [ "# @title 3. Connect to Earth Engine, or stay offline\n", "# @markdown Offline, the notebook prepares every Earth Engine request but sends none. You can\n", - "# @markdown change `USE_EARTH_ENGINE` at any time and run from the top again.\n", + "# @markdown change `RUN_UP_TO` at any time and run all again.\n", "import ee\n", "from ee.apitestcase import ApiTestCase\n", "\n", + "STAGES = (\"offline\", \"sign in\", \"upload weights\", \"export test tile\", \"compare with PyTorch\",\n", + " \"export region\")\n", + "\n", "\n", "class Stop(Exception):\n", " \"\"\"Stops the notebook with a one-line message instead of a traceback.\"\"\"\n", @@ -132,40 +141,40 @@ " return [f\"Stopped: {self}\"]\n", "\n", "\n", - "def live_step(checkbox=None, ticked=True):\n", - " \"\"\"True when an Earth Engine step should run; otherwise says why it does not.\"\"\"\n", - " if not USE_EARTH_ENGINE:\n", - " print(\"Not run: the notebook is offline (USE_EARTH_ENGINE is False in step 2).\")\n", + "def live_step(stage):\n", + " \"\"\"True when this run goes as far as `stage`; otherwise says why the step does not run.\"\"\"\n", + " if RUN_UP_TO == \"offline\":\n", + " print('Not run: the notebook is offline (RUN_UP_TO is \"offline\" in step 2).')\n", " return False\n", - " if not ticked:\n", - " print(f\"Not run: tick {checkbox} above when you are ready, then run this cell.\")\n", + " if STAGES.index(RUN_UP_TO) < STAGES.index(stage):\n", + " print(f'Not run: set RUN_UP_TO to \"{stage}\" in step 2 to go this far.')\n", " return False\n", " if not PROJECT or PROJECT == \"YOUR_EE_PROJECT\":\n", " raise Stop(\"set PROJECT to your Cloud project ID in step 2.\")\n", - " if CHECKPOINT is None:\n", - " raise Stop(\"the Earth Engine steps need your trained model: set CHECKPOINT in step 2.\")\n", " return True\n", "\n", "\n", + "if RUN_UP_TO not in STAGES:\n", + " raise Stop(f\"RUN_UP_TO must be one of: {', '.join(STAGES)}.\")\n", "# Offline, Earth Engine's own test helper stands in for a signed-in session.\n", "if \"offline\" in globals(): # undo an earlier offline run in this session\n", " offline.tearDown()\n", " ee.data._install_cloud_api_resource = _install_cloud_api_resource\n", " ee.Reset()\n", " del offline\n", - "if USE_EARTH_ENGINE:\n", - " print(\"Earth Engine mode: the next cell signs you in.\")\n", - "else:\n", + "if RUN_UP_TO == \"offline\":\n", " _install_cloud_api_resource = ee.data._install_cloud_api_resource\n", " offline = ApiTestCase()\n", " offline.setUp()\n", - " print(\"Offline mode: nothing will be sent to Earth Engine or created in a project.\")" + " print(\"Offline mode: nothing will be sent to Earth Engine or created in a project.\")\n", + "else:\n", + " print(\"Earth Engine mode: the next cell signs you in.\")" ] }, { "cell_type": "code", "execution_count": null, - "id": "70296347", + "id": "3c7c0d9f", "metadata": { "tags": [ "live" @@ -173,9 +182,9 @@ }, "outputs": [], "source": [ - "# @title Sign in (Earth Engine mode only)\n", + "# @title Sign in\n", "# @markdown The first time, follow the link with the Google account you use for Earth Engine.\n", - "if live_step():\n", + "if live_step(\"sign in\"):\n", " ee.Authenticate()\n", " ee.Initialize(project=PROJECT)\n", " print(f\"Earth Engine is ready — signed in to {PROJECT}.\")" @@ -184,14 +193,15 @@ { "cell_type": "code", "execution_count": null, - "id": "6efd6d01", + "id": "05214a7e", "metadata": {}, "outputs": [], "source": [ "# @title 4. Load the model\n", - "# @markdown Without a `CHECKPOINT`, a small demo U-Net with random weights. Your checkpoint must\n", - "# @markdown match `ENCODER`, `BANDS` and `CLASSES`: if not, this cell fails with an error listing\n", - "# @markdown the weight names that do not fit.\n", + "# @markdown - Without a `CHECKPOINT`: a small demo U-Net with random weights, which reads\n", + "# @markdown Sentinel-1 VV and VH in dB, rescaled to about −1 to 1 as (dB + 20) / 20.\n", + "# @markdown - Your checkpoint must match `ENCODER`, `BANDS` and `CLASSES`: if not, this cell fails\n", + "# @markdown with an error listing the weight names that do not fit. Its input is used as it is.\n", "import torch\n", "import segmentation_models_pytorch as smp\n", "\n", @@ -208,7 +218,9 @@ "if CHECKPOINT is None:\n", " torch.manual_seed(20261020)\n", " model = make_model(encoder_depth=3, decoder_channels=(16, 8, 4))\n", - " tile, test_tile = 64, 64 # small tiles for the demo\n", + " tile, test_tile = 64, 64 # small tiles keep the demo cheap\n", + " WEIGHTS_RECORD = \"tensor_ee_demo_weights.json\"\n", + " TEST_ASSET, REGION_ASSET = \"tensor_ee_demo_test_tile\", \"tensor_ee_demo_region\"\n", " print(\"Loaded the demo model: a small U-Net with random weights.\")\n", "else:\n", " model = make_model()\n", @@ -239,7 +251,7 @@ "\n", "\n", "def prepare(height, width):\n", - " \"\"\"A copy of the model for one input size; stops unless TensorEE can run every layer.\"\"\"\n", + " \"\"\"A copy of the model for one input size; stops unless TensorEE can build every layer.\"\"\"\n", " shape = (1, len(BANDS), height, width)\n", " adapted = adapt_smp_model(model, input_shape=shape)\n", " report = check_compatibility(adapted, shape)\n", @@ -318,38 +330,46 @@ "outputs": [], "source": [ "# @title 7. Upload the weights (once)\n", - "# @markdown - Creates the weight table assets, under `projects/PROJECT/assets/tensor_ee_weights_`,\n", - "# @markdown with one export task each, two at a time.\n", + "# @markdown - Creates the weight table assets, named `projects/PROJECT/assets/tensor_ee_weights__part_0000`\n", + "# @markdown and so on (`tensor_ee_demo_weights_…` for the demo), with one export task each, two at a time.\n", "# @markdown - Writes `WEIGHTS_RECORD`, the local list of those tables, which the prediction reads.\n", "# @markdown Keep a copy: in Colab, download it from the Files panel, since Colab's files are\n", "# @markdown lost when its runtime is reset.\n", + "# @markdown - Skipped when that file shows a finished upload; an upload that stopped part way carries on.\n", "# @markdown - Cost not measured. For scale, a trained ResNet-18 U-Net (14.3 M numbers) took 29 tables.\n", - "UPLOAD_WEIGHTS = False # @param {type:\"boolean\"}\n", - "\n", "import json\n", "import uuid\n", "from pathlib import Path\n", "\n", "from tensor_ee import export_weights\n", "\n", - "if live_step(\"UPLOAD_WEIGHTS\", UPLOAD_WEIGHTS):\n", + "if live_step(\"upload weights\"):\n", " record_path = Path(WEIGHTS_RECORD)\n", - " if record_path.exists():\n", - " raise Stop(f\"the weights are already uploaded ({record_path} exists): untick \"\n", - " \"UPLOAD_WEIGHTS and go on.\")\n", - " namespace = f\"projects/{PROJECT}/assets/tensor_ee_weights_{uuid.uuid4().hex[:12]}\"\n", - " reported = []\n", - "\n", - " def progress(record):\n", - " done = sum(shard[\"status\"] == \"COMPLETED\" for shard in record[\"shards\"])\n", - " if not reported or done != reported[-1]:\n", - " reported.append(done)\n", - " print(f\" {done} of about {tables} table(s) uploaded...\", flush=True)\n", - "\n", - " print(\"Uploading; follow it here or in the Tasks tab: https://code.earthengine.google.com/tasks\")\n", - " export_weights(parameters, namespace, record_path, on_update=progress)\n", - " print(f\"Done: the weights are uploaded, and {record_path} lists them. Untick \"\n", - " \"UPLOAD_WEIGHTS, then run on from step 8.\")" + " record = json.loads(record_path.read_text()) if record_path.exists() else None\n", + " if record and record[\"status\"] == \"ready\":\n", + " print(f\"Already uploaded: {record_path} lists the weight tables, so this step is skipped.\")\n", + " else:\n", + " if record and record[\"status\"] != \"uploading\":\n", + " raise Stop(f\"an earlier upload {record['status']} (see {record_path}). Delete that file, \"\n", + " \"and its tables in the Code Editor's Assets tab, then run again.\")\n", + " if record:\n", + " namespace = record[\"asset_namespace\"]\n", + " print(\"Carrying on with the upload that stopped part way.\", flush=True)\n", + " else:\n", + " prefix = \"tensor_ee_demo_weights\" if CHECKPOINT is None else \"tensor_ee_weights\"\n", + " namespace = f\"projects/{PROJECT}/assets/{prefix}_{uuid.uuid4().hex[:12]}\"\n", + " print(\"Uploading; follow it here or in the Tasks tab: \"\n", + " \"https://code.earthengine.google.com/tasks\", flush=True)\n", + " reported = []\n", + "\n", + " def progress(record):\n", + " done = sum(shard[\"status\"] == \"COMPLETED\" for shard in record[\"shards\"])\n", + " if not reported or done != reported[-1]:\n", + " reported.append(done)\n", + " print(f\" {done} of about {tables} table(s) uploaded...\", flush=True)\n", + "\n", + " export_weights(parameters, namespace, record_path, on_update=progress, resume=record is not None)\n", + " print(f\"Done: the weights are uploaded, and {record_path} lists them.\")" ] }, { @@ -361,7 +381,8 @@ "source": [ "# @title 8. Set up the prediction (nothing runs yet)\n", "# @markdown - Connects the model to your input image and uploaded weights, for the test tile.\n", - "# @markdown Offline, placeholder weights stand in: the prediction is complete but cannot run.\n", + "# @markdown Until the weights are uploaded, placeholder weights stand in: the prediction is\n", + "# @markdown complete but cannot run.\n", "# @markdown - Output bands: `score_k` (the model's raw scores, or logits), `p_k` (probabilities),\n", "# @markdown `class_index` and, for two classes, `binary_class` (1 where `p_1` ≥ 0.5).\n", "# @markdown - Pixels missing any input band get no prediction.\n", @@ -377,12 +398,11 @@ "\n", "\n", "def weights_record():\n", - " \"\"\"The uploaded weight tables, or offline placeholders with the same layout.\"\"\"\n", - " if USE_EARTH_ENGINE and Path(WEIGHTS_RECORD).exists():\n", - " return json.loads(Path(WEIGHTS_RECORD).read_text())\n", - " if USE_EARTH_ENGINE:\n", - " raise Stop(f\"{WEIGHTS_RECORD} not found. Upload the weights first (tick UPLOAD_WEIGHTS \"\n", - " \"in step 7), or, if you uploaded them before, put that file back in this folder.\")\n", + " \"\"\"The uploaded weight tables if the record shows a finished upload, else placeholders.\"\"\"\n", + " if RUN_UP_TO != \"offline\" and Path(WEIGHTS_RECORD).exists():\n", + " record = json.loads(Path(WEIGHTS_RECORD).read_text())\n", + " if record[\"status\"] == \"ready\":\n", + " return record\n", " return {\"format\": \"native_ee_featurecollection_weights_v2\", \"status\": \"ready\",\n", " \"layers\": {p.key: {\"shape\": list(p.shape), \"dtype\": p.dtype, \"bias\": p.bias is not None,\n", " \"source\": p.source, \"sha256\": p.sha256,\n", @@ -391,9 +411,13 @@ "\n", "\n", "def input_image():\n", - " if USE_EARTH_ENGINE:\n", - " return ee.Image(INPUT_IMAGE)\n", - " return ee.Image.constant([0] * len(BANDS)).rename(list(BANDS)) # offline stand-in\n", + " \"\"\"The model's input in Earth Engine (a blank stand-in offline).\"\"\"\n", + " if RUN_UP_TO == \"offline\":\n", + " return ee.Image.constant([0] * len(BANDS)).rename(list(BANDS))\n", + " image = ee.Image(INPUT_IMAGE).select(list(BANDS))\n", + " if CHECKPOINT is None:\n", + " image = image.add(20).divide(20) # the demo's input: dB to about -1 to 1\n", + " return image\n", "\n", "\n", "def build_prediction(plan):\n", @@ -422,14 +446,18 @@ " maxPixels=grid.width * grid.height, pyramidingPolicy={\".default\": \"sample\"}, overwrite=False)\n", "\n", "\n", + "weights_uploaded = weights_record().get(\"asset_namespace\") is not None\n", "with deep_recursion():\n", " test_prediction = build_prediction(test_plan)\n", " test_graph_bytes = len(test_prediction.serialize().encode()) # checks it can be sent\n", - "if USE_EARTH_ENGINE:\n", - " print(\"The test-tile prediction is set up with your uploaded weights. Step 9 runs it.\")\n", - "else:\n", + "if RUN_UP_TO == \"offline\":\n", " print(\"The test-tile prediction is set up offline, with placeholder weights: complete, \"\n", - " \"but it cannot run.\")" + " \"but it cannot run.\")\n", + "elif weights_uploaded:\n", + " print(\"The test-tile prediction is set up with your uploaded weights.\")\n", + "else:\n", + " print('The test-tile prediction is set up with placeholder weights; \"upload weights\" '\n", + " \"uploads the real ones.\")" ] }, { @@ -445,23 +473,31 @@ "source": [ "# @title 9. Export the test tile\n", "# @markdown - Computes the test tile in a batch export and saves it as the image asset `TEST_ASSET`\n", - "# @markdown (the tile only, not its border).\n", + "# @markdown (the tile only, not its border). Skipped when that asset already exists.\n", "# @markdown - Cost not measured. For scale, notebook 01's 224 × 224 chip used about 65 EECU-seconds.\n", - "# @markdown - If it fails for lack of memory or time, lower `TEST_TILE` in step 2 and run from the top.\n", - "EXPORT_TEST_TILE = False # @param {type:\"boolean\"}\n", - "\n", + "# @markdown - If it fails for lack of memory or time, lower `TEST_TILE` in step 2 and run all again.\n", "import time\n", "\n", "TASKS_URL = \"https://code.earthengine.google.com/tasks\"\n", + "if \"started\" not in globals():\n", + " started = {} # exports started in this session, by asset\n", + "\n", + "\n", + "def asset_exists(asset):\n", + " try:\n", + " ee.data.getAsset(asset)\n", + " return True\n", + " except ee.EEException:\n", + " return False\n", "\n", "\n", - "def wait(task, minutes=60):\n", + "def wait(task, minutes):\n", " \"\"\"Wait for an export, with a line every 5 minutes; stop if it fails.\"\"\"\n", - " started, polls = time.time(), 0\n", + " begun, polls = time.time(), 0\n", " while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", - " if time.time() - started > minutes * 60:\n", + " if time.time() - begun > minutes * 60:\n", " raise Stop(f\"still running after {minutes} minutes. It carries on in Earth Engine: \"\n", - " f\"follow it at {TASKS_URL}, and go on once it has finished.\")\n", + " f\"follow it at {TASKS_URL}, and run all again once it has finished.\")\n", " time.sleep(30)\n", " polls += 1\n", " if polls % 10 == 0:\n", @@ -471,13 +507,29 @@ " f\"(see {TASKS_URL})\")\n", "\n", "\n", - "if live_step(\"EXPORT_TEST_TILE\", EXPORT_TEST_TILE):\n", - " test_task = export_task(test_prediction, test_plan, TEST_ASSET, \"tensor_ee_test_tile\")\n", - " test_task.start()\n", - " print(f\"Export started; follow it here or at {TASKS_URL}\", flush=True)\n", - " wait(test_task)\n", - " print(f\"Done: the test tile is saved as projects/{PROJECT}/assets/{TEST_ASSET}. Untick \"\n", - " \"EXPORT_TEST_TILE, then go on to step 10.\")" + "def export_once(asset_name, make_image, plan, description, minutes):\n", + " \"\"\"Export to `asset_name` unless it exists; wait for an export this session already started.\"\"\"\n", + " asset = f\"projects/{PROJECT}/assets/{asset_name}\"\n", + " task = started.get(asset)\n", + " if task is not None and task.status()[\"state\"] not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\n", + " print(\"The export started earlier is still running; waiting for it.\", flush=True)\n", + " elif asset_exists(asset):\n", + " print(f\"Already exported: {asset} exists, so this step is skipped.\")\n", + " return\n", + " else:\n", + " if not weights_uploaded:\n", + " raise Stop('the weights are not uploaded yet: set RUN_UP_TO to \"upload weights\" or later.')\n", + " with deep_recursion():\n", + " task = export_task(make_image(), plan, asset_name, description)\n", + " task.start()\n", + " started[asset] = task\n", + " print(f\"Export started; follow it here or at {TASKS_URL}\", flush=True)\n", + " wait(task, minutes)\n", + " print(f\"Done: saved as {asset}.\")\n", + "\n", + "\n", + "if live_step(\"export test tile\"):\n", + " export_once(TEST_ASSET, lambda: test_prediction, test_plan, \"tensor_ee_test_tile\", minutes=60)" ] }, { @@ -492,14 +544,14 @@ "outputs": [], "source": [ "# @title 10. Compare the test tile with PyTorch\n", - "# @markdown - Run it once step 9 is done. It downloads the test tile's input (with its border)\n", - "# @markdown and the exported scores: 2 small requests, nothing created.\n", - "# @markdown - It then runs your model here on the same input and prints the largest difference\n", + "# @markdown - Downloads the test tile's input (with its border) and the exported scores: 2 small\n", + "# @markdown requests, nothing created.\n", + "# @markdown - Runs your model here on the same input and prints the largest difference\n", "# @markdown (target: at most 1e-4). The input must have no gaps; above 40 MB it is refused.\n", - "COMPARE_WITH_PYTORCH = False # @param {type:\"boolean\"}\n", - "\n", "import numpy as np\n", "\n", + "within_target = False\n", + "\n", "\n", "def download(image, grid, bands):\n", " names = (\"scaleX\", \"shearX\", \"translateX\", \"shearY\", \"scaleY\", \"translateY\")\n", @@ -510,12 +562,12 @@ " return np.stack([values[band] for band in bands]).astype(np.float32)\n", "\n", "\n", - "if live_step(\"COMPARE_WITH_PYTORCH\", COMPARE_WITH_PYTORCH):\n", + "if live_step(\"compare with PyTorch\"):\n", " padded = test_plan.global_grid\n", " size = padded.width * padded.height * (len(BANDS) + 1) * 4\n", " if size > 40_000_000:\n", " raise Stop(f\"the input is {size / 1e6:.0f} MB; lower TEST_TILE in step 2.\")\n", - " source = ee.Image(INPUT_IMAGE).select(list(BANDS)).toFloat().reproject(padded.projection())\n", + " source = input_image().select(list(BANDS)).toFloat().reproject(padded.projection())\n", " valid = source.mask().reduce(ee.Reducer.min()).rename(\"valid\")\n", " pixels = download(source.unmask(0).addBands(valid), padded, (*BANDS, \"valid\"))\n", " if pixels[-1].min() < 1:\n", @@ -526,11 +578,15 @@ " reference = model(torch.from_numpy(pixels[None, :-1].copy()))[0].numpy()\n", " core = reference[:, halo:halo + test_tile, halo:halo + test_tile]\n", " error = np.abs(exported - core)\n", + " within_target = bool(error.max() <= 1e-4)\n", " print(f\"Largest difference in the scores: {error.max():.3e} (target 1e-4); \"\n", " f\"mean {error.mean():.3e}; pixels with a different class: \"\n", " f\"{int((exported.argmax(0) != core.argmax(0)).sum())} of {core[0].size}.\")\n", - " print(\"Done: within the target. Earth Engine agrees with PyTorch on this tile.\" if error.max() <= 1e-4\n", - " else \"Done, but above the target: please report it as an issue, with these numbers.\")" + " print(\"Done: within the target. Earth Engine agrees with PyTorch on this tile.\" if within_target\n", + " else \"Done, but above the target: please report it as an issue, with these numbers.\")\n", + " if CHECKPOINT is None:\n", + " print(\"(The demo's random weights make a meaningless map: this shows only that Earth \"\n", + " \"Engine computes the same numbers as PyTorch.)\")" ] }, { @@ -545,19 +601,18 @@ "outputs": [], "source": [ "# @title 11. Export the region\n", - "# @markdown - Run it once step 10 is within the target. One batch export, saved as the image\n", - "# @markdown asset `REGION_ASSET` (the region only, not its border).\n", + "# @markdown - One batch export, saved as the image asset `REGION_ASSET` (the region only, not its\n", + "# @markdown border). Skipped when that asset already exists.\n", + "# @markdown - Runs only when step 10 was within the target in this run.\n", "# @markdown - Cost, memory use and time not measured at these tile sizes.\n", - "EXPORT_REGION = False # @param {type:\"boolean\"}\n", - "\n", - "if live_step(\"EXPORT_REGION\", EXPORT_REGION):\n", - " with deep_recursion():\n", - " region_task = export_task(build_prediction(region_plan), region_plan, REGION_ASSET,\n", - " \"tensor_ee_region\")\n", - " region_task.start()\n", - " print(f\"Export started; follow it here or at {TASKS_URL}\", flush=True)\n", - " wait(region_task, minutes=240)\n", - " print(f\"Done: the region is saved as projects/{PROJECT}/assets/{REGION_ASSET}.\")" + "if live_step(\"export region\"):\n", + " if not within_target:\n", + " raise Stop(\"step 10 was not within the target, so the region is not exported.\")\n", + " export_once(REGION_ASSET, lambda: build_prediction(region_plan), region_plan,\n", + " \"tensor_ee_region\", minutes=240)\n", + " if CHECKPOINT is None:\n", + " print(\"(The demo's random weights make a meaningless map: open it only to see that the \"\n", + " \"export worked.)\")" ] }, { @@ -568,9 +623,8 @@ "## Good to know\n", "\n", "- **Coming back later:** put your `WEIGHTS_RECORD` file back next to the notebook (in Colab,\n", - " upload it in the Files panel, like the checkpoint), set\n", - " `USE_EARTH_ENGINE = True`, and **Run all** with every checkbox unticked; then tick only the\n", - " step you need.\n", + " upload it in the Files panel, like the checkpoint), and **Run all**: the finished steps are\n", + " skipped.\n", "- **Evidence:** step 10 is agreement in Earth Engine for your model on one tile; keep its\n", " numbers. Accuracy is separate: evaluate the exports against your own labelled data.\n", "- **Details:** [the guide this notebook follows](https://github.com/noelEOS/tensor-ee/blob/main/docs/own-model.md)\n", diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py index caacbee..1707bea 100644 --- a/tests/test_notebooks.py +++ b/tests/test_notebooks.py @@ -8,9 +8,9 @@ - The quickstart (01) has no offline mode of its own: the test installs the client's offline signatures, and after its settings cell supplies a project name and the place and scene that its live cells would find (the preset Sun Moon Lake chip). -- 03 runs with ``USE_EARTH_ENGINE = False`` and its own offline signatures; switching it to - Earth Engine must undo them without a restart. Each of its Earth Engine steps has its own - checkbox, off as shipped. +- 03 runs with ``RUN_UP_TO = "offline"`` and its own offline signatures; switching it to + Earth Engine must undo them without a restart. Its Earth Engine steps run up to the step + ``RUN_UP_TO`` names, and skip the steps already done (checked here with stand-in tasks). - 02 is offline throughout. Notebooks that need segmentation-models-pytorch or timm skip by name without them. @@ -109,7 +109,7 @@ def test_only_live_cells_reach_earth_engine(name): def test_the_quickstart_has_one_setting_and_stops_cleanly_without_it(): sources = [_source(cell) for cell in _code("01_quickstart.ipynb")] text = "\n".join(sources) - assert "RUN_LIVE" not in text and "ApiTestCase" not in text and "restart" not in text.lower() + assert "RUN_LIVE" not in text and "ApiTestCase" not in text settings = [source for source in sources if 'PROJECT = "" # @param {type:"string"}' in source] assert len(settings) == 1 namespace = {} @@ -143,16 +143,24 @@ def status(self): assert namespace["TASKS_URL"] in str(stopped.value) -def test_live_steps_in_03_check_the_switches_and_the_project(): +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_no_notebook_mentions_a_restart(name): + # Installs use !pip, which prints no restart note; the text never asks for one. + text = "\n".join(_source(cell) for cell in _cells(name)) + assert "restart" not in text.lower() and "%pip" not in text + + +def test_live_steps_in_03_follow_one_setting_and_check_the_project(): cells = _code("03_bring_your_own_cnn.ipynb") - settings = [cell for cell in cells if re.search(r"^USE_EARTH_ENGINE = False\b", _source(cell), re.M)] + settings = [cell for cell in cells if re.search(r'^RUN_UP_TO = "offline" # @param \[', _source(cell), re.M)] assert len(settings) == 1 and f'PROJECT = "{PLACEHOLDER_PROJECT}"' in _source(settings[0]) - checkboxes = [] + options = json.loads(re.search(r"^RUN_UP_TO = .*# @param (\[.*\])$", _source(settings[0]), re.M).group(1)) + stages = [] for cell in cells: if _live(cell): - assert "live_step(" in _source(cell) - checkboxes += re.findall(r'^([A-Z_]+) = False # @param \{type:"boolean"\}$', _source(cell), re.M) - assert checkboxes == ["UPLOAD_WEIGHTS", "EXPORT_TEST_TILE", "COMPARE_WITH_PYTORCH", "EXPORT_REGION"] + stages += re.findall(r'live_step\("([^"]+)"\)', _source(cell)) + assert options == ["offline", *stages] == ["offline", "sign in", "upload weights", "export test tile", + "compare with PyTorch", "export region"] @pytest.mark.parametrize("name", sorted(EXPECTED)) @@ -230,8 +238,8 @@ def _run(name, monkeypatch, *, after_settings=None, signatures=False, then=None) raise if after_settings and "@param" in _source(cell): namespace.update(after_settings) - if "USE_EARTH_ENGINE" in namespace: - assert namespace["USE_EARTH_ENGINE"] is False + if "RUN_UP_TO" in namespace: + assert namespace["RUN_UP_TO"] == "offline" if then: then(namespace) finally: @@ -272,20 +280,77 @@ def test_bring_your_own_cnn_runs_offline_and_goes_live_without_a_restart(monkeyp import ee def switch_to_live(namespace): - # Rerunning the connect cell with USE_EARTH_ENGINE = True undoes the offline signatures. + # Rerunning the connect cell past "offline" undoes the offline signatures. connect = next((index, cell) for index, cell in enumerate(_cells("03_bring_your_own_cnn.ipynb")) if "ApiTestCase()" in _source(cell)) - namespace["USE_EARTH_ENGINE"] = True + namespace["RUN_UP_TO"] = "sign in" offline = namespace["offline"] _exec("03_bring_your_own_cnn.ipynb", *connect, namespace) assert "offline" not in namespace assert ee.data._install_cloud_api_resource is namespace["_install_cloud_api_resource"] assert ee.data.getAlgorithms is offline.old_get_algorithms assert ee.data.getMapId is offline.old_get_map_id - namespace["USE_EARTH_ENGINE"] = False + namespace["RUN_UP_TO"] = "offline" namespace = _run("03_bring_your_own_cnn.ipynb", monkeypatch, then=switch_to_live) assert namespace["CHECKPOINT"] is None assert (namespace["halo"], namespace["alignment"]) == (112, 8) assert len(namespace["region_plan"].tiles) == 2 and len(namespace["test_plan"].tiles) == 1 assert namespace["test_graph_bytes"] > 0 + + +def _live_cell(name, title): + return next((index, cell) for index, cell in enumerate(_cells(name)) + if _live(cell) and _source(cell).startswith(title)) + + +def test_bring_your_own_cnn_skips_the_steps_already_done(monkeypatch, tmp_path, capsys): + import tensor_ee + name = "03_bring_your_own_cnn.ipynb" + + class Task: + def __init__(self, states): + self.states, self.id = list(states), "stand-in" + + def start(self): + starts.append(self) + + def status(self): + return {"state": self.states.pop(0) if len(self.states) > 1 else self.states[0]} + + uploads, starts, existing = [], [], set() + + def check(namespace): + record = tmp_path / "weights.json" + namespace.update(RUN_UP_TO="export region", PROJECT="offline-check", WEIGHTS_RECORD=str(record)) + monkeypatch.setattr(tensor_ee, "export_weights", lambda *args, **kwargs: uploads.append(kwargs)) + upload = _live_cell(name, "# @title 7.") + _exec(name, *upload, namespace) # no record: a new upload + record.write_text(json.dumps({"status": "uploading", "asset_namespace": "stand-in"})) + _exec(name, *upload, namespace) # stopped part way: carried on + record.write_text(json.dumps({"status": "ready"})) + _exec(name, *upload, namespace) # finished: skipped + assert [kwargs["resume"] for kwargs in uploads] == [False, True] + assert "Already uploaded" in capsys.readouterr().out + record.write_text(json.dumps({"status": "failed"})) + with pytest.raises(Exception) as stopped: + _exec(name, *upload, namespace) + assert type(stopped.value).__name__ == "Stop" + + namespace["RUN_UP_TO"] = "offline" + _exec(name, *_live_cell(name, "# @title 9."), namespace) # defines the export helpers only + namespace.update(RUN_UP_TO="export region", weights_uploaded=True, + asset_exists=lambda asset: asset in existing, + export_task=lambda image, plan, asset_name, description: Task(["RUNNING", "COMPLETED"])) + monkeypatch.setattr(namespace["time"], "sleep", lambda seconds: None) + export_once = namespace["export_once"] + export_once("tile", lambda: None, None, "test", minutes=60) # a new export + existing.add("projects/offline-check/assets/tile") + export_once("tile", lambda: None, None, "test", minutes=60) # it exists: skipped + assert len(starts) == 1 and "Already exported" in capsys.readouterr().out + namespace["started"]["projects/offline-check/assets/region"] = Task(["RUNNING", "RUNNING", "COMPLETED"]) + export_once("region", lambda: None, None, "region", minutes=60) # still running: waited for + assert len(starts) == 1 and "still running; waiting" in capsys.readouterr().out + namespace["RUN_UP_TO"] = "offline" + + _run(name, monkeypatch, then=check) From b6298aa64ca0a8d0d970cc9cd5ee4d284a04ff05 Mon Sep 17 00:00:00 2001 From: Noel Ivan Ulloa <127597880+noelEOS@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:19:39 +0800 Subject: [PATCH 9/9] Notebook 03: a progress bar and the expected time for the weight upload --- notebooks/03_bring_your_own_cnn.ipynb | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/notebooks/03_bring_your_own_cnn.ipynb b/notebooks/03_bring_your_own_cnn.ipynb index a95ef92..36594ab 100644 --- a/notebooks/03_bring_your_own_cnn.ipynb +++ b/notebooks/03_bring_your_own_cnn.ipynb @@ -330,6 +330,9 @@ "outputs": [], "source": [ "# @title 7. Upload the weights (once)\n", + "# @markdown - **This takes a while: about 7 minutes for the demo, longer for a bigger model.**\n", + "# @markdown The bar below updates every 15 seconds. Earth Engine reports each table only as\n", + "# @markdown waiting, running or done, so the bar moves one table at a time.\n", "# @markdown - Creates the weight table assets, named `projects/PROJECT/assets/tensor_ee_weights__part_0000`\n", "# @markdown and so on (`tensor_ee_demo_weights_…` for the demo), with one export task each, two at a time.\n", "# @markdown - Writes `WEIGHTS_RECORD`, the local list of those tables, which the prediction reads.\n", @@ -338,6 +341,7 @@ "# @markdown - Skipped when that file shows a finished upload; an upload that stopped part way carries on.\n", "# @markdown - Cost not measured. For scale, a trained ResNet-18 U-Net (14.3 M numbers) took 29 tables.\n", "import json\n", + "import time\n", "import uuid\n", "from pathlib import Path\n", "\n", @@ -360,15 +364,19 @@ " namespace = f\"projects/{PROJECT}/assets/{prefix}_{uuid.uuid4().hex[:12]}\"\n", " print(\"Uploading; follow it here or in the Tasks tab: \"\n", " \"https://code.earthengine.google.com/tasks\", flush=True)\n", - " reported = []\n", + " upload_began = time.time()\n", "\n", " def progress(record):\n", - " done = sum(shard[\"status\"] == \"COMPLETED\" for shard in record[\"shards\"])\n", - " if not reported or done != reported[-1]:\n", - " reported.append(done)\n", - " print(f\" {done} of about {tables} table(s) uploaded...\", flush=True)\n", + " states = [shard[\"status\"] for shard in record[\"shards\"]]\n", + " done, running = states.count(\"COMPLETED\"), states.count(\"RUNNING\")\n", + " total = max(tables, len(states))\n", + " minutes, seconds = divmod(int(time.time() - upload_began), 60)\n", + " line = (f\" [{'█' * done}{'░' * (total - done)}] {done} of about {total} tables done, \"\n", + " f\"{running} running, {minutes} min {seconds:02d} s\")\n", + " print(line.ljust(80), end=\"\\r\", flush=True)\n", "\n", " export_weights(parameters, namespace, record_path, on_update=progress, resume=record is not None)\n", + " print()\n", " print(f\"Done: the weights are uploaded, and {record_path} lists them.\")" ] },