Estimating canopy height with RasterFlow

This notebook will guide you through estimating canopy height from aerial imagery, using Wherobots RasterFlow and the Meta and World Resources Institute’s Canopy Height prediction model (Meta CHM v1)¹. This is an open source regression model that can predict the height of the tree canopy from high resolution imagery. The result of the model is a raster where each pixel is the estimated tree canopy height in meters.

It was trained on high-resolution imagery (Maxar Vivid2 mosaic with 0.5m resolution), existing labeled canopy height maps and airborne laser scans (LIDAR).

Running inference on high resolution imagery

The Meta CHM v1 model was trained on 50cm resolution data, so we will select imagery with a similar resolution for our testing.

The National Agriculture Imagery Program (NAIP) provides aerial imagery for the United States, capturing high-resolution images during the agricultural growing seasons. The NAIP dataset contains a mixture of resolutions (30cm, 60cm and 1m).

Preview: model inputs and outputs

The interactive map linked below shows the input imagery and model outputs for this notebook’s example AOI. Toggle layers in the sidebar to compare the input mosaic and the raster model output side-by-side.

Layers (top layer first, as in the map’s layer list):

CHM model outputs: Single-band raster where each pixel is the estimated tree canopy height in meters
CHM input mosaic: RGB NAIP aerial imagery the model runs on

View the interactive map here.

Selecting an Area of Interest (AOI)

We will choose an Area of Interest (AOI) for our analysis.

The area around Nashua, NH has a combination of urban settings, parks and forests so we will try out the model there.

import shapely
import wkls
from sedona.spark import *

config = SedonaContext.builder().getOrCreate()
sedona = SedonaContext.create(config)

# Generate a WKT geometry for Nashua, NH using Well-Known Locations (https://github.com/wherobots/wkls)
aoi = wkls.us.nh.nashua.wkt()

min_lon, min_lat, max_lon, max_lat = shapely.from_wkt(aoi).bounds
print(f"Bounds: [{min_lon:.4f}, {min_lat:.4f}, {max_lon:.4f}, {max_lat:.4f}]")

Selecting a time range and verifying NAIP coverage

NAIP is flown state by state on a multi-year cycle, and the Meta CHM v1 recipe sources 60cm imagery. The most recent 60cm imagery for New Hampshire was captured in 2021, so we will set our time range accordingly. If you change the AOI or the time range, the 60cm imagery the model needs may not exist. The USDA’s NAIP coverage map (PDF, 2002-2025) shows what is available where.

from datetime import datetime

# Date range for imagery to be used by the model
start_date = datetime(2021, 1, 1)
end_date = datetime(2022, 1, 1)

# The Meta CHM v1 recipe runs on 60cm NAIP imagery from the RGB (visual) product. Wherobots Open Data
# indexes every scene of that product with its footprint, resolution (gsd) and acquisition time.
MODEL_RES = 0.6
sedona.createDataFrame([(aoi,)], ["wkt"]).selectExpr("ST_GeomFromWKT(wkt) AS geom").createOrReplaceTempView("aoi")

# Scenes at the model's resolution touching the AOI (the table and wkls are both EPSG:4326)
sedona.sql(f"""
    SELECT n.geometry, n.year, n.datetime
    FROM wherobots_open_data.imagery.usda_naip_visual_stac n
    JOIN aoi a ON ST_Intersects(n.geometry, a.geom)
    WHERE n.gsd = {MODEL_RES}
""").cache().createOrReplaceTempView("nearby")

# The recipe needs the AOI fully inside the union of the scene footprints in the date range
coverage = sedona.sql(f"""
    SELECT count(*) AS scenes, array_sort(collect_set(n.year)) AS years,
           ST_Contains(ST_Union_Agg(n.geometry), first(a.geom)) AS covered
    FROM nearby n CROSS JOIN aoi a
    WHERE n.datetime >= TIMESTAMP '{start_date.isoformat()}Z' AND n.datetime < TIMESTAMP '{end_date.isoformat()}Z'
""").first()

if not coverage["covered"]:
    # Years that would work, applying the same full-coverage test to each year on its own
    available = [
        row.year
        for row in sedona.sql("""
            SELECT n.year FROM nearby n CROSS JOIN aoi a
            GROUP BY n.year
            HAVING ST_Contains(ST_Union_Agg(n.geometry), first(a.geom))
            ORDER BY n.year
        """).collect()
    ]
    raise ValueError(
        f"{MODEL_RES:g}m NAIP does not fully cover this AOI between "
        f"{start_date:%Y-%m-%d} and {end_date:%Y-%m-%d}. "
        + (
            f"Years with complete {MODEL_RES:g}m coverage over this AOI: {available}."
            if available
            else f"No year has complete {MODEL_RES:g}m NAIP coverage over this AOI."
        )
    )

print(f"NAIP coverage is available. Found {coverage['scenes']} NAIP scenes at {MODEL_RES:g}m covering the AOI")
print(f"Acquisition years: {coverage['years']}")

Using the Wherobots Tasks API

Each RasterFlow workflow runs as a task through the Wherobots Python SDK. run_task below starts a task run, waits for it to finish, and returns its output value. Recipes return an object with input_mosaic and result; standalone tasks return a URI. Call wb.tasks.get("<task name>").input_schema to see every input a task accepts.

import time

from wherobots import TaskRunState, Wherobots

# In a Wherobots notebook the SDK uses your Wherobots credentials, elsewhere it reads WHEROBOTS_API_KEY
wb = Wherobots()


def run_task(task_name, inputs, poll_seconds=30):
    """Run a RasterFlow task, wait for it to finish, and return its output value."""
    run = wb.task_runs.get(wb.tasks.run(task_name, inputs, region="aws-us-west-2").task_run_id)
    print(f"Started {task_name} run {run.task_run_id}")
    while run.is_active:
        time.sleep(poll_seconds)
        run = wb.task_runs.get(run.task_run_id)
    if run.state != TaskRunState.SUCCEEDED:
        raise RuntimeError(f"{task_name} run {run.task_run_id} ended {run.state}: {run.error_message}")
    # A successful run can temporarily have no results if the output read fails.
    for _ in range(3):
        if run.results is not None:
            break
        time.sleep(poll_seconds)
        run = wb.task_runs.get(run.task_run_id)
    if run.results is None:
        raise RuntimeError(
            f"Results for successful run {run.task_run_id} are unavailable. "
            "Fetch this run again with wb.task_runs.get(run_id); do not resubmit the task."
        )
    # The single task output can be a URI or a structured recipe result.
    (output,) = run.results.values()
    return output

Running a model

RasterFlow has pre-defined workflows to simplify orchestration of the processing steps for model inference.These steps include:

Ingesting imagery for the specified Area of Interest (AOI)
Generating a seamless image from multiple image tiles (a mosaic)
Running inference with the selected model

The recipe returns input_mosaic for the source imagery and result for the predictions. Each contains a mosaic_index_uri for its GeoParquet index and a first_row_mosaic for its first Zarr store.

Note: This step will take approximately 20 minutes to complete.

recipe_output = run_task(
    "predict_mosaic_recipe",
    {
        # Area of Interest as WKT (EPSG:4326)
        "aoi": aoi,

        # Date range for imagery to be used by the model (set in the coverage check above), as RFC 3339 timestamps
        "start": start_date.isoformat() + "Z",
        "end": end_date.isoformat() + "Z",

        # Coordinate Reference System for the output mosaic
        "target_crs": "EPSG:3857",

        # The model recipe to be used for inference (CHM in this case)
        "model_recipe": "meta-chm-v1",
    },
)

# The recipe returns both the input imagery and the prediction mosaic.
input_mosaic_index_uri = recipe_output["input_mosaic"]["mosaic_index_uri"]
input_mosaic_store = recipe_output["input_mosaic"]["first_row_mosaic"]
print(f"Input imagery: {input_mosaic_store}")

model_output_index_uri = recipe_output["result"]["mosaic_index_uri"]

# Inspect the prediction index (one row per output mosaic location)
model_output_index_df = sedona.read.format("geoparquet").load(model_output_index_uri)
model_output_index_df.show(truncate=60)

# This example AOI has one geometry, so we select the first (only) mosaic location.
model_output_store = model_output_index_df.first()["location"]

model_output_store

(Optional) Build an optimized Zarr for visualization

RasterFlow writes its outputs as Zarr stores at native resolution. To explore them interactively on cloud.wherobots.com/map, you can build an optimized multiscale Zarr. build_zarr_multiscales adds downsampled overview levels (image pyramids) to the store so the map can stream coarse tiles when zoomed out and full-resolution pixels when zoomed in.

This step is optional and can take a few minutes for large outputs, so the code below is commented out by default — uncomment it to run it.

Pass aoi to select a chunk-aligned bounding window from a single Zarr source. It accepts polygon or multipolygon WKT in EPSG:4326, or an S3 GeoJSON or GeoParquet URI with a CRS. Pixels outside the polygon remain within that rectangular window. Omit aoi to process the full store.

# optimized_store = run_task("build_zarr_multiscales", {"source_store": model_output_store, "aoi": aoi})
# optimized_store

Visualizing outputs

You can visualize the Zarr, GeoParquet, and other geospatial outputs using cloud.wherobots.com/map.

References

Tolan, J., Yang, H.-I., Nosarzewski, B., Couairon, G., Vo, H. V., Brandt, J., Spore, J., Majumdar, S., Haziza, D., Vamaraju, J., et al. (2024). Very high resolution canopy height maps from RGB imagery using self-supervised vision transformer and convolutional decoder trained on aerial lidar. Remote Sensing of Environment, 300, 113888.