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..10456fd --- /dev/null +++ b/docs/start-here.md @@ -0,0 +1,81 @@ +# 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. +- **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 + +- **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, 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. 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). + +## 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:** 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 + +- **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 + ([validation](validation.md)). +- **Accuracy:** not measured by TensorEE; measure it on your own labelled data. + +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 new file mode 100644 index 0000000..a9505b2 --- /dev/null +++ b/notebooks/01_quickstart.ipynb @@ -0,0 +1,308 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2ac26326", + "metadata": {}, + "source": [ + "# 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", + "- **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": "code", + "id": "7361375f", + "metadata": {}, + "source": [ + "# @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", + "import sys\n", + "from pathlib import Path\n", + "\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))\n", + "print(\"TensorEE is installed.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "33d242a2", + "metadata": {}, + "source": [ + "# @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", + "import time\n", + "\n", + "ASSET_NAME = \"tensor_ee_quickstart_\" + time.strftime(\"%Y%m%d_%H%M%S\") # a new asset each run\n", + "\n", + "\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", + "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.\")\n", + "print(f\"Settings saved: project {PROJECT}, point ({LONGITUDE}, {LATITUDE}), \"\n", + " f\"scenes from {START_DATE} to {END_DATE}.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "9b2cdbb8", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "# @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", + "ee.Authenticate()\n", + "ee.Initialize(project=PROJECT)\n", + "print(f\"Earth Engine is ready — signed in to {PROJECT}.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "376fa9ae", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "# @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", + "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(f\"Found the chip ({CHIP} × {CHIP} px in {crs}) and a Sentinel-1 scene that covers it: {scene_id}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "c035f536", + "metadata": {}, + "source": [ + "# @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 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", + "\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", + "prediction = (ee.ImageCollection.fromImages([chip]).mosaic().select(BANDS)\n", + " .setDefaultProjection(grid.projection()).clip(grid.geometry())\n", + " .updateMask(grid.mask()).toFloat())\n", + "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": [] + }, + { + "cell_type": "code", + "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", + "# @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(\"Export started; it usually takes a few minutes. Follow it here or at\", TASKS_URL)\n", + "wait_for(task)\n", + "print(\"Done: the water map is saved as\", asset_id)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "d32406d8", + "metadata": { + "tags": [ + "live" + ] + }, + "source": [ + "# @title 7. See the radar image and the water map\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", + " 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", + "centre_lon, centre_lat = grid.geometry().centroid(1).transform(\"EPSG:4326\", 1).coordinates().getInfo()\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=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", + "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": "ee640c02", + "metadata": {}, + "source": [ + "## Good to know\n", + "\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", + "- **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." + ] + } + ], + "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..7b7d107 --- /dev/null +++ b/notebooks/02_is_my_model_supported.ipynb @@ -0,0 +1,399 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "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 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.\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, and attention models such as SegFormer." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ec064c72", + "metadata": {}, + "outputs": [], + "source": [ + "# @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", + "execution_count": null, + "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", + " \"\"\"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: 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 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", + " 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 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 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": "af9bb1a2", + "metadata": {}, + "source": [ + "## 3. A supported model: a U-Net with a ResNet-18 encoder\n", + "\n", + "- 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", + "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", + "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))" + ] + }, + { + "cell_type": "markdown", + "id": "65fcf0bc", + "metadata": {}, + "source": [ + "- **\"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": "af182bec", + "metadata": {}, + "source": [ + "## 4. Refused models\n", + "\n", + "### A LinkNet: transposed convolutions\n", + "\n", + "- 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", + "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))" + ] + }, + { + "cell_type": "markdown", + "id": "f593c0f3", + "metadata": {}, + "source": [ + "### DeepLabV3+: dilated convolutions, and more\n", + "\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", + "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))" + ] + }, + { + "cell_type": "markdown", + "id": "975e2f88", + "metadata": {}, + "source": [ + "**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 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": "25e7a7d4", + "metadata": {}, + "source": [ + "## 5. Attention models: the converter\n", + "\n", + "- 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", + "execution_count": null, + "id": "515d11e5", + "metadata": {}, + "outputs": [], + "source": [ + "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)" + ] + }, + { + "cell_type": "markdown", + "id": "23fd38cf", + "metadata": {}, + "source": [ + "- 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", + "execution_count": null, + "id": "5121948a", + "metadata": {}, + "outputs": [], + "source": [ + "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])" + ] + }, + { + "cell_type": "markdown", + "id": "fffc365e", + "metadata": {}, + "source": [ + "- 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", + "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())" + ] + }, + { + "cell_type": "markdown", + "id": "fb434a3d", + "metadata": {}, + "source": [ + "- 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", + "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])" + ] + }, + { + "cell_type": "markdown", + "id": "b34a78b5", + "metadata": {}, + "source": [ + "## 6. Your model\n", + "\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", + "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 tries the CNN check first (SMP models and plain PyTorch CNNs), then the converter." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7c46d932", + "metadata": {}, + "outputs": [], + "source": [ + "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(\"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)" + ] + }, + { + "cell_type": "markdown", + "id": "9020a2f7", + "metadata": {}, + "source": [ + "## What next\n", + "\n", + "- **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) 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 checked settings." + ] + } + ], + "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..36594ab --- /dev/null +++ b/notebooks/03_bring_your_own_cnn.ipynb @@ -0,0 +1,660 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "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", + "Runs a U-Net you trained with\n", + "[segmentation-models-pytorch](https://github.com/qubvel-org/segmentation_models.pytorch) (SMP)\n", + "inside Earth Engine, on your own imagery.\n", + "\n", + "**How to use it**\n", + "\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", + "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", + "execution_count": null, + "id": "c4686e47", + "metadata": {}, + "outputs": [], + "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", + "print(\"TensorEE is installed.\")" + ] + }, + { + "cell_type": "markdown", + "id": "ca2d2707", + "metadata": {}, + "source": [ + "## 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 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)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27953840", + "metadata": {}, + "outputs": [], + "source": [ + "# 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", + "\n", + "# Your model: must match how it was trained\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 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", + "\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 (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\n", + "REGION_ASSET = \"my_unet_region\" # image asset for the region\n", + "\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: {model_name}. This run goes up to \"{RUN_UP_TO}\" in project {PROJECT}.')" + ] + }, + { + "cell_type": "code", + "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 `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", + "\n", + " def _render_traceback_(self):\n", + " return [f\"Stopped: {self}\"]\n", + "\n", + "\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 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", + " 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 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.\")\n", + "else:\n", + " print(\"Earth Engine mode: the next cell signs you in.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3c7c0d9f", + "metadata": { + "tags": [ + "live" + ] + }, + "outputs": [], + "source": [ + "# @title Sign in\n", + "# @markdown The first time, follow the link with the Google account you use for Earth Engine.\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}.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "05214a7e", + "metadata": {}, + "outputs": [], + "source": [ + "# @title 4. Load the model\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", + "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 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", + " model.load_state_dict(torch.load(CHECKPOINT, map_location=\"cpu\", weights_only=True), strict=True)\n", + " tile, test_tile = TILE, TEST_TILE\n", + " print(f\"Loaded your model from {CHECKPOINT}.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4e6dbb6c", + "metadata": {}, + "outputs": [], + "source": [ + "# @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", + " \"\"\"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", + " if report.status != \"supported\":\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", + "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", + "execution_count": null, + "id": "79dd01f2", + "metadata": {}, + "outputs": [], + "source": [ + "# @title 6. Plan the tiles\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), (\"border\", halo)):\n", + " if value % alignment:\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", + " return plan_tiles(padded, side, halo, alignment=alignment, roi=region)\n", + "\n", + "\n", + "def describe(plan, name):\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 {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\")\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", + "execution_count": null, + "id": "ad3edbbe", + "metadata": { + "tags": [ + "live" + ] + }, + "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", + "# @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", + "import json\n", + "import time\n", + "import uuid\n", + "from pathlib import Path\n", + "\n", + "from tensor_ee import export_weights\n", + "\n", + "if live_step(\"upload weights\"):\n", + " record_path = Path(WEIGHTS_RECORD)\n", + " 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", + " upload_began = time.time()\n", + "\n", + " def progress(record):\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.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "efe612f2", + "metadata": {}, + "outputs": [], + "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 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", + "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_record():\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", + " \"assets\": [\"projects/YOUR_EE_PROJECT/assets/weights_not_uploaded\"]}\n", + " for p in parameters}}\n", + "\n", + "\n", + "def input_image():\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", + " 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 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", + " 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", + "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 RUN_UP_TO == \"offline\":\n", + " print(\"The test-tile prediction is set up offline, with placeholder weights: complete, \"\n", + " \"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.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "78c0ce07", + "metadata": { + "tags": [ + "live" + ] + }, + "outputs": [], + "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). 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 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):\n", + " \"\"\"Wait for an export, with a line every 5 minutes; stop if it fails.\"\"\"\n", + " begun, polls = time.time(), 0\n", + " while (state := task.status()[\"state\"]) not in (\"COMPLETED\", \"FAILED\", \"CANCELLED\"):\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 run all again 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", + " f\"(see {TASKS_URL})\")\n", + "\n", + "\n", + "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)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4b544941", + "metadata": { + "tags": [ + "live" + ] + }, + "outputs": [], + "source": [ + "# @title 10. Compare the test tile with PyTorch\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", + "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", + " 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 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 = 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 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", + " 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 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.)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "020c43c6", + "metadata": { + "tags": [ + "live" + ] + }, + "outputs": [], + "source": [ + "# @title 11. Export the region\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", + "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.)\")" + ] + }, + { + "cell_type": "markdown", + "id": "3dee0517", + "metadata": {}, + "source": [ + "## 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), 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", + " 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." + ] + } + ], + "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..1707bea --- /dev/null +++ b/tests/test_notebooks.py @@ -0,0 +1,356 @@ +"""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, 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_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. +""" +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|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\"'`)]*") +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): + return json.loads((NOTEBOOKS / name).read_text())["cells"] + + +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"): + 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_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 + 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_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) + + +@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'^RUN_UP_TO = "offline" # @param \[', _source(cell), re.M)] + assert len(settings) == 1 and f'PROJECT = "{PLACEHOLDER_PROJECT}"' in _source(settings[0]) + 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): + 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)) +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 in (PLACEHOLDER_PROJECT, ""), value + + +@contextlib.contextmanager +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 + + 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) + 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: 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 _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(signatures): + try: + for index, cell in enumerate(_cells(name)): + if cell["cell_type"] != "code" or _live(cell): + continue + 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_UP_TO" in namespace: + assert namespace["RUN_UP_TO"] == "offline" + if then: + then(namespace) + 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 + 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 found["sha256"] == 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_and_goes_live_without_a_restart(monkeypatch): + import ee + + def switch_to_live(namespace): + # 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["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["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)