Skip to content

Commit 3b1ad43

Browse files
wietzesuijkerwietzesuijker
authored
feat: Add Sentinel-2 RGB to GeoZarr workflow with embedded STAC block (#13)
1 parent 319b1df commit 3b1ad43

File tree

1 file changed

+193
-0
lines changed

1 file changed

+193
-0
lines changed

‎docs/examples/06_embedded_STAC_block.ipynb

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "201df5fa",
6+
"metadata": {},
7+
"source": [
8+
"# Create a STAC-aware GeoZarr RGB tile of a Sentinel-2 scene over Vienna.\n",
9+
"\n",
10+
"This example demonstrates a best-practice workflow for creating a STAC-aware,\n",
11+
"interoperable GeoZarr store from Sentinel-2 Cloud-Optimized GeoTIFFs (COGs).\n",
12+
"\n",
13+
"The key steps are:\n",
14+
"1. Query Earth-Search for a recent, low-cloud Sentinel-2 L2A scene.\n",
15+
"2. Stream and stack the source RGB bands into a single, lazy xarray.DataArray.\n",
16+
"3. Prepare a GeoZarr-compliant xarray.Dataset in memory, which includes\n",
17+
" embedding STAC metadata and refining the CRS attributes.\n",
18+
"4. Write the final, consolidated GeoZarr store.\n",
19+
"5. Verify the result by reopening the store and inspecting the STAC metadata.\n",
20+
"6. Display a coarsened preview of the RGB image."
21+
]
22+
},
23+
{
24+
"cell_type": "code",
25+
"execution_count": null,
26+
"id": "be2dcbf1",
27+
"metadata": {},
28+
"outputs": [],
29+
"source": [
30+
"\"\"\"Create a STAC-aware GeoZarr RGB tile of a Sentinel-2 scene over Vienna.\"\"\"\n",
31+
"\n",
32+
"import shutil\n",
33+
"from datetime import date\n",
34+
"from pathlib import Path\n",
35+
"\n",
36+
"import jsonschema\n",
37+
"import matplotlib.pyplot as plt\n",
38+
"import pystac_client\n",
39+
"import rioxarray as rxr\n",
40+
"import xarray as xr\n",
41+
"\n",
42+
"# 1. Query Earth-Search for a recent, low-cloud Sentinel-2 L2A scene.\n",
43+
"print(\"Step 1: Querying for a Sentinel-2 scene over Vienna...\")\n",
44+
"API = \"https://earth-search.aws.element84.com/v1\"\n",
45+
"coll = \"sentinel-2-l2a\"\n",
46+
"bbox = [16.20, 48.10, 16.45, 48.30] # Vienna\n",
47+
"today = date.today()\n",
48+
"last_year = today.replace(year=today.year - 1)\n",
49+
"daterange = f\"{last_year:%Y-%m-%d}/{today:%Y-%m-%d}\"\n",
50+
"\n",
51+
"item = next(\n",
52+
" pystac_client.Client.open(API)\n",
53+
" .search(\n",
54+
" collections=[coll],\n",
55+
" bbox=bbox,\n",
56+
" datetime=daterange,\n",
57+
" query={\"eo:cloud_cover\": {\"lt\": 5}},\n",
58+
" limit=1,\n",
59+
" )\n",
60+
" .items(),\n",
61+
" None,\n",
62+
")\n",
63+
"if not item:\n",
64+
" raise RuntimeError(\"No Sentinel-2 scene found for the specified criteria.\")\n",
65+
"print(f\"Found scene: {item.id} (Cloud cover: {item.properties['eo:cloud_cover']:.2f}%)\")\n",
66+
"\n",
67+
"\n",
68+
"# 2. Stream and stack the source RGB bands into a lazy DataArray.\n",
69+
"print(\"\\nStep 2: Streaming and stacking source COG bands...\")\n",
70+
"bands = [\"red\", \"green\", \"blue\"]\n",
71+
"rgb = xr.concat(\n",
72+
" [\n",
73+
" rxr.open_rasterio(\n",
74+
" item.assets[b].href, chunks={\"band\": 1, \"x\": 2048, \"y\": 2048}, masked=True\n",
75+
" ).assign_coords(band=[b])\n",
76+
" for b in bands\n",
77+
" ],\n",
78+
" dim=\"band\",\n",
79+
")\n",
80+
"rgb.name = \"radiance\"\n",
81+
"# Assign the Coordinate Reference System from the STAC item's properties.\n",
82+
"rgb = rgb.rio.write_crs(item.properties[\"proj:code\"])\n",
83+
"\n",
84+
"\n",
85+
"# 3. Prepare a complete, GeoZarr-compliant xarray.Dataset in memory.\n",
86+
"print(\"\\nStep 3: Preparing the final xarray.Dataset with all metadata...\")\n",
87+
"# Convert the stacked RGB DataArray into a Dataset and name the data variable\n",
88+
"radiance_ds = rgb.to_dataset(name=\"radiance\")\n",
89+
"\n",
90+
"# Grab the original GeoTransform from the DataArray\n",
91+
"transform = rgb.rio.transform()\n",
92+
"\n",
93+
"# Ensure x/y are recognized as spatial dims, then write transform + CRS at the dataset level\n",
94+
"radiance_ds = (\n",
95+
" radiance_ds.rio.set_spatial_dims(x_dim=\"x\", y_dim=\"y\")\n",
96+
" .rio.write_transform(transform)\n",
97+
" .rio.write_crs(item.properties[\"proj:code\"])\n",
98+
")\n",
99+
"\n",
100+
"# Build the STAC Item metadata dictionary\n",
101+
"gsd = min(item.assets[b].to_dict().get(\"gsd\", 10) for b in bands)\n",
102+
"store_path = Path(f\"{coll}_{'_'.join(bands)}_{item.id}.zarr\")\n",
103+
"stac_metadata = {\n",
104+
" \"type\": \"Item\",\n",
105+
" \"stac_version\": \"1.0.0\",\n",
106+
" \"id\": item.id,\n",
107+
" \"bbox\": item.bbox,\n",
108+
" \"geometry\": item.geometry,\n",
109+
" \"properties\": {\n",
110+
" \"datetime\": item.properties[\"datetime\"],\n",
111+
" \"proj:code\": item.properties[\"proj:code\"],\n",
112+
" \"platform\": item.properties[\"platform\"],\n",
113+
" \"instruments\": item.properties[\"instruments\"],\n",
114+
" \"eo:cloud_cover\": item.properties[\"eo:cloud_cover\"],\n",
115+
" \"gsd\": gsd,\n",
116+
" },\n",
117+
" \"assets\": {\n",
118+
" \"data\": {\n",
119+
" \"href\": store_path.name,\n",
120+
" \"type\": \"application/x-zarr\",\n",
121+
" \"roles\": [\"data\"],\n",
122+
" }\n",
123+
" },\n",
124+
" \"license\": item.properties.get(\"license\", \"proprietary\"),\n",
125+
"}\n",
126+
"# Basic STAC schema check\n",
127+
"jsonschema.validate(\n",
128+
" instance=stac_metadata,\n",
129+
" schema={\"type\": \"object\", \"required\": [\"type\", \"id\", \"stac_version\", \"assets\"]},\n",
130+
")\n",
131+
"# Embed it in the dataset attrs\n",
132+
"radiance_ds.attrs[\"stac\"] = stac_metadata\n",
133+
"\n",
134+
"\n",
135+
"# 4. Write the final, consolidated GeoZarr store to disk.\n",
136+
"print(f\"\\nStep 4: Writing the final GeoZarr store to {store_path.resolve()}...\")\n",
137+
"if store_path.exists():\n",
138+
" shutil.rmtree(store_path)\n",
139+
"\n",
140+
"radiance_ds.chunk({\"y\": 512, \"x\": 512}).to_zarr(store_path, mode=\"w\", consolidated=True)\n",
141+
"\n",
142+
"\n",
143+
"# 5. Verify the result by reopening the store and checking for STAC metadata.\n",
144+
"print(\"\\nStep 5: Verifying the created GeoZarr store...\")\n",
145+
"with xr.open_zarr(store_path, consolidated=True) as verified_ds:\n",
146+
" if \"stac\" not in verified_ds.attrs:\n",
147+
" raise RuntimeError(\"FAIL: STAC metadata was not found.\")\n",
148+
" print(\"✅ SUCCESS: Embedded STAC metadata was found.\")\n",
149+
"\n",
150+
"\n",
151+
"# 6. Display a coarsened preview of the RGB image.\n",
152+
"print(\"\\nStep 6: Generating and displaying a coarsened preview...\")\n",
153+
"with xr.open_zarr(store_path, consolidated=True) as ds_to_plot:\n",
154+
" preview = ds_to_plot.radiance.coarsen(y=8, x=8, boundary=\"trim\").mean()\n",
155+
"\n",
156+
" # Apply a simple contrast stretch for better visualization.\n",
157+
" p2, p98 = preview.quantile(0.02), preview.quantile(0.98)\n",
158+
"\n",
159+
" preview_stretched = preview.clip(min=p2, max=p98)\n",
160+
" preview_stretched = (preview_stretched - p2) / (p98 - p2)\n",
161+
"\n",
162+
" rgb_preview = preview_stretched.transpose(\"y\", \"x\", \"band\").values\n",
163+
"\n",
164+
" fig, ax = plt.subplots(figsize=(8, 8))\n",
165+
" ax.imshow(rgb_preview)\n",
166+
" ax.set_title(f\"Coarsened Preview of {item.id}\")\n",
167+
" ax.set_axis_off()\n",
168+
" plt.show()"
169+
]
170+
}
171+
],
172+
"metadata": {
173+
"kernelspec": {
174+
"display_name": "mm",
175+
"language": "python",
176+
"name": "python3"
177+
},
178+
"language_info": {
179+
"codemirror_mode": {
180+
"name": "ipython",
181+
"version": 3
182+
},
183+
"file_extension": ".py",
184+
"mimetype": "text/x-python",
185+
"name": "python",
186+
"nbconvert_exporter": "python",
187+
"pygments_lexer": "ipython3",
188+
"version": "3.12.11"
189+
}
190+
},
191+
"nbformat": 4,
192+
"nbformat_minor": 5
193+
}

0 commit comments

Comments
 (0)