"""Shared data foundation for the "AI habitat classification with Python" series (Peak District, UK). Every part of the series imports this module so that all parts use the same study area, 10 m grid, class scheme, labels, Sentinel-2/Sentinel-1 composites and spatial train/test split. import sys; sys.path.insert(0, "../ai-habitat-classification-common") import habitat_common as hc See README.md in this folder for the full interface. """ from __future__ import annotations import json import time import warnings from pathlib import Path import numpy as np import pandas as pd HERE = Path(__file__).resolve().parent DATA = HERE / "data" # downloads and cached composites (gitignored) SAMPLES_CSV = HERE / "samples" / "habitat_samples.csv" # committed: the shared sample design BLOCKS_CSV = HERE / "samples" / "blocks.csv" # committed: block_id -> split and fold LABELS_GPKG = DATA / "living_england_aoi.gpkg" LABELS_TIF = DATA / "labels_10m.tif" # --- Study area and grid ---------------------------------------------------------------------------- CRS = "EPSG:27700" # British National Grid RES = 10 # metres AOI_BOUNDS = (400000, 377500, 425000, 400000) # xmin, ymin, xmax, ymax: Dark Peak, Edale and Hope valleys AOI_NAME = "Dark Peak (Kinder Scout, Bleaklow, Edale and the Hope Valley)" YEAR = 2023 # imagery year: matches the Living England 2022-23 epoch BLOCK_SIZE = 2500 # spatial blocks for the train/test split, metres BLOCK_BUFFER = 50 # no samples within this distance of a block edge, metres SEED = 42 # --- Class scheme ----------------------------------------------------------------------------------- # id: value in label rasters and samples (0 = no label). le: Living England 2022-23 "Prmry_H" values. CLASSES = [ {"id": 1, "key": "bog", "name": "Blanket bog", "colour": "#7b3294", "le": ["Bog"]}, {"id": 2, "key": "heath", "name": "Heath", "colour": "#c2a5cf", "le": ["Dwarf Shrub Heath"]}, {"id": 3, "key": "bracken", "name": "Bracken", "colour": "#e08214", "le": ["Bracken"]}, {"id": 4, "key": "seminat_grass", "name": "Semi-natural grassland", "colour": "#fee08b", "le": ["Unimproved Grassland"]}, # acid on the gritstone, calcareous on the limestone near Castleton {"id": 5, "key": "improved_grass", "name": "Improved grassland", "colour": "#a6d96a", "le": ["Improved and Semi-Improved Grassland", "Arable and Horticultural"]}, {"id": 6, "key": "broadleaf", "name": "Broadleaf woodland", "colour": "#1a9850", "le": ["Broadleaved, Mixed and Yew Woodland", "Scrub"]}, {"id": 7, "key": "conifer", "name": "Conifer woodland", "colour": "#00441b", "le": ["Coniferous Woodland"]}, {"id": 8, "key": "water", "name": "Water", "colour": "#2166ac", "le": ["Water"]}, {"id": 9, "key": "built_bare", "name": "Built-up and bare", "colour": "#636363", "le": ["Built-up Areas and Gardens", "Bare Ground", "Bare Sand"]}, ] EXCLUDED_LE = ["Fen, Marsh and Swamp"] # too rare and ambiguous here: left unlabelled CLASS_IDS = [c["id"] for c in CLASSES] CLASS_NAMES = {c["id"]: c["name"] for c in CLASSES} CLASS_KEYS = {c["id"]: c["key"] for c in CLASSES} LE_TO_ID = {le: c["id"] for c in CLASSES for le in c["le"]} RELIABLE = ["High", "Very High"] # Living England reliability values used for sampling # --- Sentinel-2 / Sentinel-1 ----------------------------------------------------------------------- STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1" S2_BANDS = ["B02", "B03", "B04", "B05", "B06", "B07", "B08", "B8A", "B11", "B12"] SCL_CLEAR = [4, 5, 6, 7] # vegetation, bare soil, water, unclassified SCL_CLOUD = [3, 8, 9, 10] # cloud shadow, cloud medium/high probability, thin cirrus (dilated by 20 m) SEASONS = {"winter": [1, 2, 12], "spring": [3, 4, 5], "summer": [6, 7, 8], "autumn": [9, 10, 11]} LE_SERVICE = ("https://services.arcgis.com/JJzESW51TqeY9uat/ArcGIS/rest/services/" "Living_England_2022_2023/FeatureServer") LE_LAYERS = [3, 5, 6] # biogeographic zones that intersect the AOI def geobox(): """The shared 10 m GeoBox (odc-geo) in EPSG:27700: 2,500 columns x 2,250 rows.""" from odc.geo.geobox import GeoBox return GeoBox.from_bbox(AOI_BOUNDS, crs=CRS, resolution=RES) def transform(): """Affine transform of the shared grid (rasterio convention, north-up).""" from rasterio.transform import from_origin return from_origin(AOI_BOUNDS[0], AOI_BOUNDS[3], RES, RES) def shape(): """(rows, cols) of the shared grid.""" return (int((AOI_BOUNDS[3] - AOI_BOUNDS[1]) / RES), int((AOI_BOUNDS[2] - AOI_BOUNDS[0]) / RES)) def aoi_bbox_wgs84(): """AOI bounding box in longitude/latitude, for STAC searches.""" from pyproj import Transformer t = Transformer.from_crs(CRS, "EPSG:4326", always_xy=True) xs, ys = t.transform([AOI_BOUNDS[0], AOI_BOUNDS[0], AOI_BOUNDS[2], AOI_BOUNDS[2]], [AOI_BOUNDS[1], AOI_BOUNDS[3], AOI_BOUNDS[1], AOI_BOUNDS[3]]) return (min(xs), min(ys), max(xs), max(ys)) def class_cmap(): """(ListedColormap, BoundaryNorm) for label/prediction rasters with values 1..9 (mask 0 before plotting).""" from matplotlib.colors import BoundaryNorm, ListedColormap cmap = ListedColormap([c["colour"] for c in CLASSES]) cmap.set_bad(alpha=0) return cmap, BoundaryNorm(np.arange(0.5, len(CLASSES) + 1.5), cmap.N) # --- Labels ------------------------------------------------------------------------------------------ def _get_json(url, params, tries=5): import requests for i in range(tries): try: r = requests.post(url, data=params, timeout=180) # POST: long objectIds lists overflow a GET URL r.raise_for_status() js = r.json() if "error" in js: raise RuntimeError(js["error"]) return js except Exception: if i == tries - 1: raise time.sleep(5 * (i + 1)) def download_living_england(force=False, verbose=True): """Download Living England 2022-23 polygons that intersect the AOI (Natural England, OGL v3). Pages through the public ArcGIS FeatureServer (no login), clips to the AOI and saves a GeoPackage. Returns a GeoDataFrame in EPSG:27700 with columns le_segid, le_habitat, le_reliability, le_source, class_id. """ import geopandas as gpd from shapely.geometry import box if LABELS_GPKG.exists() and not force: return gpd.read_file(LABELS_GPKG) DATA.mkdir(exist_ok=True) env = ",".join(str(v) for v in AOI_BOUNDS) frames = [] for layer in LE_LAYERS: url = f"{LE_SERVICE}/{layer}/query" ids = _get_json(url, {"where": "1=1", "geometry": env, "geometryType": "esriGeometryEnvelope", "inSR": 27700, "spatialRel": "esriSpatialRelIntersects", "returnIdsOnly": "true", "f": "json"})["objectIds"] or [] ids = sorted(ids) for i in range(0, len(ids), 1000): js = _get_json(url, {"objectIds": ",".join(map(str, ids[i:i + 1000])), "outFields": "SegID,Prmry_H,Relblty,Source", "outSR": 27700, "f": "geojson"}) frames.append(gpd.GeoDataFrame.from_features(js["features"], crs=CRS)) if verbose: print(f"Living England 2022-23, zone layer {layer}: {len(ids):,} polygons intersect the AOI") gdf = gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), geometry="geometry", crs=CRS) gdf = gdf.rename(columns={"Prmry_H": "le_habitat", "Relblty": "le_reliability", "Source": "le_source", "SegID": "le_segid"}) gdf["geometry"] = gdf.geometry.make_valid().intersection(box(*AOI_BOUNDS)) gdf = gdf[~gdf.geometry.is_empty].explode(index_parts=False) gdf = gdf[gdf.geometry.geom_type == "Polygon"] gdf["class_id"] = gdf["le_habitat"].map(LE_TO_ID).fillna(0).astype("uint8") gdf = gdf[["le_segid", "le_habitat", "le_reliability", "le_source", "class_id", "geometry"]].reset_index(drop=True) gdf.to_file(LABELS_GPKG, driver="GPKG") return gdf def load_labels(): """Living England polygons for the AOI (downloads on first use).""" return download_living_england() def rasterize_labels(gdf=None, reliable_only=False): """Rasterise class_id onto the shared grid (uint8, 0 = unlabelled). Pixel-centre rule.""" from rasterio.features import rasterize gdf = load_labels() if gdf is None else gdf if reliable_only: gdf = gdf[gdf["le_reliability"].isin(RELIABLE)] gdf = gdf[gdf["class_id"] > 0] return rasterize(zip(gdf.geometry, gdf["class_id"]), out_shape=shape(), transform=transform(), fill=0, dtype="uint8") def label_raster(): """Full Living England label map on the shared grid (all reliabilities), cached as data/labels_10m.tif.""" import rasterio if LABELS_TIF.exists(): with rasterio.open(LABELS_TIF) as src: return src.read(1) arr = rasterize_labels() write_geotiff(LABELS_TIF, arr, nodata=0) return arr def write_geotiff(path, arr, nodata=None, dtype=None): """Write a 2-D (rows, cols) or 3-D (bands, rows, cols) array on the shared grid as a compressed GeoTIFF.""" import rasterio arr = np.asarray(arr) if arr.ndim == 2: arr = arr[None] dtype = dtype or arr.dtype Path(path).parent.mkdir(parents=True, exist_ok=True) with rasterio.open(path, "w", driver="GTiff", height=arr.shape[1], width=arr.shape[2], count=arr.shape[0], dtype=dtype, crs=CRS, transform=transform(), nodata=nodata, compress="deflate", tiled=True, blockxsize=512, blockysize=512) as dst: dst.write(arr.astype(dtype)) # --- Samples and the spatial split ------------------------------------------------------------------- def block_id(x, y): """Spatial block index (row-major from the north-west corner) for BNG coordinates.""" ncols = int(np.ceil((AOI_BOUNDS[2] - AOI_BOUNDS[0]) / BLOCK_SIZE)) bc = ((np.asarray(x) - AOI_BOUNDS[0]) // BLOCK_SIZE).astype(int) br = ((AOI_BOUNDS[3] - np.asarray(y)) // BLOCK_SIZE).astype(int) return br * ncols + bc def load_blocks(): """GeoDataFrame of the spatial blocks: block_id, split ("train"/"test"), fold, geometry (samples/blocks.csv).""" import geopandas as gpd from shapely.geometry import box df = pd.read_csv(BLOCKS_CSV) ncols = int(np.ceil((AOI_BOUNDS[2] - AOI_BOUNDS[0]) / BLOCK_SIZE)) geoms = [] for b in df["block_id"]: r, c = divmod(int(b), ncols) x0, y1 = AOI_BOUNDS[0] + c * BLOCK_SIZE, AOI_BOUNDS[3] - r * BLOCK_SIZE geoms.append(box(x0, max(y1 - BLOCK_SIZE, AOI_BOUNDS[1]), min(x0 + BLOCK_SIZE, AOI_BOUNDS[2]), y1)) return gpd.GeoDataFrame(df, geometry=geoms, crs=CRS) def load_samples(split=None): """The committed sample design (samples/habitat_samples.csv) as a DataFrame. Columns: sample_id, x, y (BNG pixel centre), row, col (indices into the shared grid), class_id, class_key, le_habitat, le_reliability, le_source, block_id, split ("train" or "test"), fold (0-4 for train, -1 for test). """ df = pd.read_csv(SAMPLES_CSV) return df if split is None else df[df["split"] == split].reset_index(drop=True) def sample_values(arr, samples): """Values of a (rows, cols) or (bands, rows, cols) array at each sample's row/col.""" arr = np.asarray(arr) r, c = samples["row"].to_numpy(), samples["col"].to_numpy() return arr[r, c] if arr.ndim == 2 else arr[:, r, c].T # --- Sentinel-2 -------------------------------------------------------------------------------------- def stac_client(): import planetary_computer import pystac_client return pystac_client.Client.open(STAC_URL, modifier=planetary_computer.sign_inplace) def search_s2(start=f"{YEAR}-01-01", end=f"{YEAR}-12-31", max_cloud=90): """Sentinel-2 L2A items over the AOI, one per date and tile (newest processing baseline kept).""" items = stac_client().search(collections=["sentinel-2-l2a"], bbox=aoi_bbox_wgs84(), datetime=f"{start}/{end}", query={"eo:cloud_cover": {"lt": max_cloud}}).item_collection() best = {} for it in items: key = (it.datetime.date(), it.properties["s2:mgrs_tile"]) if key not in best or it.properties["s2:processing_baseline"] > best[key].properties["s2:processing_baseline"]: best[key] = it return sorted(best.values(), key=lambda i: i.datetime) def load_s2(items, bands=S2_BANDS, chunks=None): """Lazy Sentinel-2 cube on the shared grid: dask-backed xarray.Dataset of raw uint16 DNs plus SCL. Reflectance = (DN - offset) / 10000, with the per-date offset from s2_offsets(). 20 m bands are resampled bilinearly to the 10 m grid; SCL with nearest neighbour. """ import odc.stac odc.stac.configure_rio(cloud_defaults=True) ds = odc.stac.load(items, bands=list(bands) + ["SCL"], geobox=geobox(), groupby="solar_day", chunks=chunks or {"time": 1, "x": 2500, "y": 2250}, dtype="uint16", nodata=0, resampling={"*": "bilinear", "SCL": "nearest"}) return ds.assign_coords(time=ds.time.dt.floor("D")) def s2_offsets(items): """Per-date BOA offset (1000 for processing baseline >= 04.00, else 0) as a pandas Series.""" return pd.Series({pd.Timestamp(it.datetime.date()): 1000 if float(it.properties["s2:processing_baseline"]) >= 4.0 else 0 for it in items}) def s2_clear_masks(items=None, cache=True): """Clear-sky masks for every date: (bool DataArray time x y x x, clear-fraction Series). clear = SCL in SCL_CLEAR and not within 20 m of SCL_CLOUD. Cached as data/s2_clear_.npz. """ import xarray as xr from scipy import ndimage items = items or search_s2() f = DATA / f"s2_clear_{items[0].datetime.year}.npz" if cache and f.exists(): z = np.load(f) clear = xr.DataArray(np.unpackbits(z["bits"], axis=0, count=z["times"].size).astype(bool), dims=("time", "y", "x"), coords={"time": pd.to_datetime(z["times"])}, name="clear") else: scl = load_s2(items, bands=[])["SCL"].compute() cloudy = np.isin(scl.values, SCL_CLOUD) grown = ndimage.binary_dilation(cloudy, iterations=2, axes=(1, 2)) clear = xr.DataArray(np.isin(scl.values, SCL_CLEAR) & ~grown, dims=("time", "y", "x"), coords={"time": scl.time.values}, name="clear") DATA.mkdir(exist_ok=True) np.savez_compressed(f, bits=np.packbits(clear.values, axis=0), times=clear.time.values.astype("datetime64[D]")) frac = clear.mean(("y", "x")).to_series() return clear, frac def s2_composite(name, months, year=YEAR, bands=S2_BANDS, min_clear=0.05, cache=True, verbose=True): """Cloud-masked median surface reflectance for the given months of one year on the shared grid. Returns an xarray.Dataset of float32 reflectance (0-1), one variable per band, NaN where no clear view exists, plus "n_clear" (number of clear dates per pixel). Cached as data/s2__.tif (int16 reflectance x 10000, nodata -9999; last band n_clear). """ import rasterio import xarray as xr f = DATA / f"s2_{year}_{name}.tif" if cache and f.exists(): with rasterio.open(f) as src: arr = src.read().astype("float32") ds = xr.Dataset({b: (("y", "x"), np.where(arr[i] == -9999, np.nan, arr[i] / 10000)) for i, b in enumerate(bands)}) ds["n_clear"] = (("y", "x"), arr[-1]) return ds import dask items = search_s2(f"{year}-01-01", f"{year}-12-31") clear, frac = s2_clear_masks(items, cache=cache) times = frac.index[(frac >= min_clear) & frac.index.month.isin(months)] t0 = time.time() ds = load_s2(items, bands=bands).sel(time=times) off = xr.DataArray(s2_offsets(items).reindex(times).to_numpy(), dims="time", coords={"time": ds.time}) refl = ((ds[list(bands)].astype("float32") - off) / 10000).where(clear.sel(time=times).values) with dask.config.set(scheduler="threads", num_workers=16), warnings.catch_warnings(): warnings.simplefilter("ignore", RuntimeWarning) med = refl.chunk({"time": -1, "y": 450, "x": 500}).median("time").compute() n = clear.sel(time=times).sum("time").values.astype("float32") stack = [np.where(np.isnan(med[b].values), -9999, np.round(med[b].values * 10000)) for b in bands] write_geotiff(f, np.stack(stack + [n]), nodata=-9999, dtype="int16") if verbose: print(f"S2 {name} {year}: {len(times)} dates with >= {min_clear:.0%} clear sky, built in {time.time() - t0:.0f} s") out = xr.Dataset({b: (("y", "x"), med[b].values.astype("float32")) for b in bands}) out["n_clear"] = (("y", "x"), n) return out def s2_seasonal_composites(year=YEAR, **kw): """Dict season -> s2_composite(...) for SEASONS (winter = Jan, Feb and Dec of the same year).""" return {s: s2_composite(s, m, year=year, **kw) for s, m in SEASONS.items()} def s2_dates(year=YEAR, min_clear=0.05): """Clear fraction of the AOI for each date in the year with at least min_clear clear sky (pandas Series).""" items = search_s2(f"{year}-01-01", f"{year}-12-31") _, frac = s2_clear_masks(items) return frac[frac >= min_clear] def s2_sample_timeseries(samples=None, year=YEAR, bands=S2_BANDS, min_clear=0.05, cache=True, verbose=True): """Per-date Sentinel-2 reflectance at every sample pixel (for time-series models). Returns a long DataFrame: sample_id, date, clear (bool), then one float32 reflectance column per band (NaN where not clear). Covers every date with >= min_clear of the AOI clear (50 dates in 2023). Cached as data/s2__samples_timeseries.parquet; built for the full committed sample set. """ f = DATA / f"s2_{year}_samples_timeseries.parquet" if cache and f.exists(): ts = pd.read_parquet(f) else: allsamp = load_samples() items = search_s2(f"{year}-01-01", f"{year}-12-31") clear, frac = s2_clear_masks(items, cache=cache) times = frac.index[frac >= min_clear] off = s2_offsets(items) ds = load_s2(items, bands=bands).sel(time=times) r, c = allsamp["row"].to_numpy(), allsamp["col"].to_numpy() t0, parts = time.time(), [] for i, t in enumerate(times): day = ds.sel(time=t).compute() ok = clear.sel(time=t).values[r, c] part = pd.DataFrame({"sample_id": allsamp["sample_id"].to_numpy(), "date": t, "clear": ok}) for b in bands: v = (day[b].values[r, c].astype("float32") - off[t]) / 10000 part[b] = np.where(ok, v, np.nan).astype("float32") parts.append(part) if verbose: print(f" S2 samples {i + 1}/{len(times)} {t.date()} ({time.time() - t0:.0f} s)", flush=True) ts = pd.concat(parts, ignore_index=True) DATA.mkdir(exist_ok=True) ts.to_parquet(f, index=False) if samples is not None: ts = ts[ts["sample_id"].isin(samples["sample_id"])].reset_index(drop=True) return ts # --- Sentinel-1 RTC ---------------------------------------------------------------------------------- def search_s1(start=f"{YEAR}-01-01", end=f"{YEAR}-12-31", orbit_state=None): """Sentinel-1 RTC (gamma0, radiometrically terrain corrected) items over the AOI.""" q = {"sat:orbit_state": {"eq": orbit_state}} if orbit_state else None items = stac_client().search(collections=["sentinel-1-rtc"], bbox=aoi_bbox_wgs84(), datetime=f"{start}/{end}", query=q).item_collection() return sorted(items, key=lambda i: i.datetime) def load_s1(items, chunks=None): """Lazy Sentinel-1 RTC cube (vv, vh, linear gamma0, float32) on the shared grid, grouped by solar day.""" import odc.stac odc.stac.configure_rio(cloud_defaults=True) ds = odc.stac.load(items, bands=["vv", "vh"], geobox=geobox(), groupby="solar_day", chunks=chunks or {"time": 1, "x": 2500, "y": 2250}, dtype="float32", nodata=np.nan, resampling="bilinear") return ds.assign_coords(time=ds.time.dt.floor("D")) def s1_composite(name, months, year=YEAR, orbit_state=None, cache=True, verbose=True): """Median VV and VH gamma0 for the given months of one year on the shared grid. Returns an xarray.Dataset with vv, vh (linear), vv_db, vh_db and ratio_db (VH - VV in dB). Cached as data/s1__[_].tif (float32 linear vv, vh). """ import rasterio import xarray as xr f = DATA / f"s1_{year}_{name}{'_' + orbit_state if orbit_state else ''}.tif" if cache and f.exists(): with rasterio.open(f) as src: vv, vh = src.read().astype("float32") else: import dask t0 = time.time() items = [i for i in search_s1(f"{year}-01-01", f"{year}-12-31", orbit_state) if i.datetime.month in months] ds = load_s1(items) with dask.config.set(scheduler="threads", num_workers=16), warnings.catch_warnings(): warnings.simplefilter("ignore", RuntimeWarning) med = ds.where(ds > 0).chunk({"time": -1, "y": 450, "x": 500}).median("time").compute() vv, vh = med["vv"].values.astype("float32"), med["vh"].values.astype("float32") write_geotiff(f, np.stack([vv, vh]), nodata=np.nan, dtype="float32") if verbose: print(f"S1 {name} {year}: {ds.sizes['time']} dates, built in {time.time() - t0:.0f} s") with np.errstate(divide="ignore", invalid="ignore"): out = xr.Dataset({"vv": (("y", "x"), vv), "vh": (("y", "x"), vh), "vv_db": (("y", "x"), 10 * np.log10(vv)), "vh_db": (("y", "x"), 10 * np.log10(vh))}) out["ratio_db"] = out["vh_db"] - out["vv_db"] return out def describe(): """A short summary of the shared configuration.""" return {"crs": CRS, "res": RES, "aoi_bounds": AOI_BOUNDS, "shape": shape(), "year": YEAR, "block_size": BLOCK_SIZE, "classes": {c["id"]: c["name"] for c in CLASSES}} if __name__ == "__main__": print(json.dumps(describe(), indent=2))