Building Sentinel-2 mosaics with RasterFlow

Build a harvest-season median, an all-bands median over a custom window, and a closest-valid composite for a single reference date. Each example uses the same Area of Interest (AOI).

Choosing a Sentinel-2 composite

s2med_planting and s2med_harvest use latitude-based seasonal windows. The harvest example contains RGB, NIR, and a valid-pixel count:

• s2med_harvest:B02 – Blue
• s2med_harvest:B03 – Green
• s2med_harvest:B04 – Red
• s2med_harvest:B08 – NIR
• s2med_harvest:N_VALID_PIXELS – Number of valid observations used in the median

For a custom UTC date window, choose one of these datasets:

• s2_median produces one per-pixel RGB+NIR median over [start, end), labelled by start.
• s2_median_all_bands uses the same window and returns all 12 L2A reflectance bands.
• s2_closest_valid produces one composite for an explicit target_date, using the inclusive candidate window [start, end].
• s2_closest_valid_before_after produces two closest-valid composites at the window endpoints.

Setup and Imports

import shapely
import wkls
import xarray as xr
from sedona.spark import *

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

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 the URI of its output. Call wb.tasks.get("<task name>").input_schema to see every input a task accepts.

import time

import shapely

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 URI."""
    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."
        )
    # Each RasterFlow task has one output, named by the task (e.g. mosaic_index_uri)
    (output_uri,) = run.results.values()
    return output_uri

Selecting an Area of Interest (AOI)

We will use Haskell County, Kansas as our AOI. This is an agricultural region in the US Great Plains, ideal for demonstrating harvest-season imagery.

# Generate a WKT geometry for Haskell County, Kansas using Well-Known Locations (https://github.com/wherobots/wkls)
aoi = wkls.us.ks.haskellcounty.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}]")

Building the harvest season mosaic

This step builds a Sentinel-2 median composite mosaic for the harvest season only.

The workflow will:

1. Query Sentinel-2 imagery for the AOI and date range
2. Filter to harvest season dates based on latitude
3. Apply cloud and quality filtering — excluding scenes with >= 75% cloud cover and masking pixels flagged as cloud, cloud shadow, cirrus, saturated/defective, or unclassified in the SCL band
4. Compute a pixel-wise median composite across remaining valid observations
5. Output the result as a Zarr store

Note: This step will take approximately 10-15 minutes to complete.If you want to skip this build_mosaics step, you can uncomment the next cell and start with a pre-generated mosaic.

# Use the pre-generated mosaic instead
# mosaic_store = "s3://wherobots-examples/rasterflow/mosaics/haskell_harvest.zarr"

mosaic_index_uri = run_task(
    "build_mosaics",
    {
        # Dataset type for harvest season median mosaic
        "datasets": ["s2med_harvest"],

        # Area of Interest
        "aoi": aoi,

        # Date range for imagery to be used
        "start": "2024-01-01T00:00:00Z",
        "end": "2025-01-01T00:00:00Z",

        # Output CRS (Web Mercator)
        "target_crs": "EPSG:3857",
    },
)

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

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

mosaic_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": mosaic_store, "aoi": aoi})
# The same call works on the all-bands store: {"source_store": all_bands_store, "aoi": aoi}
# optimized_store

Visualizing outputs

You can visualize the Zarr outputs using cloud.wherobots.com/map.

Building an all-bands median over a custom window

Use s2_median_all_bands when you need a specific date window and bands beyond RGB+NIR. It returns B01, B02, B03, B04, B05, B06, B07, B08, B8A, B09, B11, and B12 at 10 m. B10 is absent from L2A.

This example creates one time step over UTC [2024-07-01, 2024-09-01), labelled 2024-07-01. It uses no seasonal heuristic. Cloud masking and per-pixel median calculation follow the harvest example above. Band names use the s2_median_all_bands prefix.

all_bands_index_uri = run_task(
    "build_mosaics",
    {
        "datasets": ["s2_median_all_bands"],
        "aoi": aoi,
        "start": "2024-07-01T00:00:00Z",
        "end": "2024-09-01T00:00:00Z",
        "target_crs": "EPSG:3857",
    },
)

all_bands_store = sedona.read.format("geoparquet").load(all_bands_index_uri).first()["location"]
print(f"All-bands mosaic store: {all_bands_store}")

Building one closest-valid composite

s2_closest_valid requires a reference date separate from the candidate window. This example searches UTC [2021-12-01, 2022-02-01] and produces one time step labelled 2022-01-01. The reference must fall within the window, and start must be earlier than end. Pass RFC 3339 timestamps; a Z suffix means UTC.

For each pixel, the compositor chooses the nearest SCL-valid observation among the selected scenes. Earlier acquisitions win equal-distance ties. Source dates may differ across pixels; the output label is the reference date. This method uses validity and date proximity, without ranking atmospheric quality. Its reflectance bands use the s2_closest_valid prefix.

Pass target_date only with s2_closest_valid. Other datasets and mixed requests reject it. For two endpoint composites, use s2_closest_valid_before_after without target_date.

closest_valid_index_uri = run_task(
    "build_mosaics",
    {
        "datasets": ["s2_closest_valid"],
        "aoi": aoi,
        "start": "2021-12-01T00:00:00Z",
        "end": "2022-02-01T00:00:00Z",
        "target_date": "2022-01-01T00:00:00Z",
        "target_crs": "EPSG:3857",
    },
)

closest_valid_store = sedona.read.format("geoparquet").load(closest_valid_index_uri).first()["location"]
print(f"Closest-valid mosaic store: {closest_valid_store}")

Next steps

• Compare with a planting-season median using s2med_planting.
• Run the FTW model for field boundary detection in RasterFlow_FTW.ipynb.
• Use the mosaic for model inference with the run_mosaic_inference task.

Set storage_uri to a persistent storage prefix when you need to retain outputs. Temporary user storage expires after seven days.