# Install AerEO and any required plugins for this notebook (Google Colab)
!pip install -q "aereo[viz]"
Step-by-step Sentinel-2 extraction¶
This notebook runs the full AerEO pipeline — Search, Build tasks and Extract (read and write) — step by step.
Each of these functions can be replaced by ANY function that conforms to the inputs and outputs of each step as follows:
| Step | Input | Output |
|---|---|---|
| Search | (collections, intersects, start_datetime, end_datetime, **kwargs) | GeoDataFrame[AssetSchema] |
| Build | (GeoDataFrame[AssetSchema], ExtractionJob, **kwargs) | Sequence[ExtractionTask] |
| Read | (task: ExtractionTask, **kwargs) | xr.Dataset |
| Write | (ds: xr.Dataset, path: str, **kwargs) | str |
No YAML config¶
This notebook builds the ExtractionJob directly in Python, so there is no separate YAML config file to copy. The equivalent YAML config is shown in the Sentinel-2 (nir, red) tutorial.
Configuring AOI, Grid and Extraction pipeline¶
Before we start, we must configure our ExtractionJob. This means defining the grid to use, the aoi and the pipeline functions to use for each step.
If we know all of this beforehand, we can run our pipeline anywhere by serializing the job, and define different pipelines using only config files!
from aereo.builtins.search import search_stac
from datetime import datetime, timezone
from shapely.geometry import Polygon
from aereo.pipeline import ExtractionJob
from aereo.builtins import read_odc_stac, write_geotiff
# AOI polygon — Chocón reservoir, Argentina (inlined so the notebook has no file dependencies).
aoi_polygon = Polygon(
[
(-68.90986824592407, -39.23705421799603),
(-68.65925870907353, -39.23705421799603),
(-68.65925870907353, -39.41589522092947),
(-68.90986824592407, -39.41589522092947),
(-68.90986824592407, -39.23705421799603),
]
)
grid_dist = 10_000
job = ExtractionJob(
name="sentinel2_sample",
grid_dist=grid_dist,
output_uri="/tmp/aereo_extraction",
read=read_odc_stac,
write=write_geotiff,
target_aoi=aoi_polygon,
)
Step 1 — Search: search_stac¶
The search provider queries the STAC API and returns a GeoDataFrame of matched assets. Each row corresponds to one requested asset (e.g. red, nir) from one STAC item. All parameters are passed explicitly.
assets = search_stac(
stac_api_url="https://earth-search.aws.element84.com/v1",
collections={"sentinel-2-l2a": ["red", "nir"]},
intersects=aoi_polygon,
start_datetime=datetime(2024, 1, 1, tzinfo=timezone.utc),
end_datetime=datetime(2024, 1, 10, tzinfo=timezone.utc),
)
print(f"\u2713 Found {len(assets)} asset rows")
print("Columns:", list(assets.columns))
print("First rows:")
assets[["id", "collection", "channel_id", "crs", "href"]].head()
Step 2 — Build tasks: build_grouped_tasks¶
The task builder takes the search-result GeoDataFrame and the ExtractionJob and produces a list of ExtractionTask objects. Tasks are grouped by start time and native CRS, and chunked by cells_per_task.
from aereo.builtins.task_builder import build_grouped_tasks
tasks = build_grouped_tasks(
search_results=assets,
job=job,
)
print(f"\u2713 Built {len(tasks)} extraction task(s)")
for i, task in enumerate(tasks):
print(f"Task {i}: {task}")
# Inspect the first task in detail
task = tasks[0]
print("Grid cells:", task.grid_cells)
print("Task context:", dict(task.task_context))
print(f"Number of asset rows in task: {len(task.assets)}")
# Extraction stages available on the task (delegated from the job)
print("Extraction stages:")
print(" read callable:", callable(task.job.read))
print(" write callable:", callable(task.job.write))
Step 3 — Read: read_odc_stac¶
The reader reconstructs pystac.Item objects from the asset table and uses odc.stac.load to build a lazy xarray.Dataset in the native CRS of the STAC items.
ds_native = job.read(task)
print("Native dataset:")
ds_native
# ds_native["red"].plot()
Step 4 — Write: write_geotiff¶
The writer serialises a dataset to a GeoTIFF at the path supplied by the caller and returns the written path.
from pathlib import Path
# Write the native dataset to a temporary path
out_path = Path(job.output_uri) / "demo_native_red.tif"
# crate the parent directory if it doesn't exist
out_path.parent.mkdir(parents=True, exist_ok=True)
written = job.write(ds=ds_native, path=out_path)
print(f"\u2713 Wrote {written}")
Batch execution with an executor¶
To run many tasks in parallel, pass the task list to an Executor. For COG/IO-bound extractors like this one, LocalExecutor(use_threads=True) is the safer choice in Jupyter and avoids the pickling issues that can make ProcessPoolExecutor hang.
Partial functions¶
Because the executor calls the read and write functions for each task in this case, it is important that, if extra params are needed, these functions are passed as partial functions. This is done in the ExtractionJob constructor, and the executor simply calls them with the task as the only argument. Usually, this is done automatically by Hydra when the job is constructed from a config file, but it can also be done manually as shown below.
from aereo.executors import LocalExecutor
# set job with partial functions for read and write
from functools import partial
# create a write partial function passing rioxarray/to_raster kwargs
write_partial = partial(write_geotiff, compress="LZW", tags={"creator": "aereo"})
job = ExtractionJob(
name="sentinel2_sample",
grid_dist=grid_dist,
output_uri="/tmp/aereo_extraction",
read=read_odc_stac,
write=write_partial,
target_aoi=aoi_polygon,
)
# lets build again the tasks with the new job
tasks = build_grouped_tasks(
search_results=assets,
job=job,
)
executor = LocalExecutor(workers=4, use_threads=True)
artifacts = executor(tasks)
print(f"\u2713 Extracted {len(artifacts)} artifact row(s) from {len(tasks)} task(s)")
artifacts[["id", "uri", "grid_cell", "start_time"]].head()
# lets check that chunks are ok
import rioxarray
rioxarray.open_rasterio(artifacts["uri"].iloc[0])
from aereo.viz import plot_artifact_patches
plot_artifact_patches(artifacts, ds_factor=1, cmap="viridis")
