"""Deep learning for satellite imagery, part 1: build a labelled training dataset from Sentinel-2. Builds a cloud-free Sentinel-2 L2A summer composite of Nottingham (UK) from the Microsoft Planetary Computer, pairs it with two real label sources, ESA WorldCover 2021 (land cover) and Overture Maps building footprints, cuts georeferenced 128 x 128 pixel chips, splits them into train/validation/test sets by spatial block (and shows why a random split leaks), checks the class balance and saves everything as GeoTIFFs plus a GeoPackage index that part 2 (train-unet-satellite-segmentation-pytorch) loads. Run with: python make_training_chips.py """ import platform import time from importlib.metadata import version from pathlib import Path import dask import geopandas as gpd import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import odc.stac import pandas as pd import planetary_computer import pystac_client import rasterio import rasterio.features import xarray as xr from matplotlib.colors import ListedColormap from matplotlib.patches import Patch, Rectangle from odc.geo.geobox import GeoBox from overturemaps import core as overture from pyproj import Transformer from rasterio.transform import from_origin from scipy import ndimage from shapely.geometry import box HERE = Path(__file__).parent OUT, DATA = HERE / "output", HERE / "data" CHIPS = DATA / "chips" for d in (OUT, DATA, CHIPS / "images", CHIPS / "masks"): d.mkdir(parents=True, exist_ok=True) plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9}) pd.set_option("display.width", 140) dask.config.set(scheduler="threads", num_workers=8) print("=== versions ===") print(f"python {platform.python_version()} | odc-stac {odc.stac.__version__} | pystac-client {pystac_client.__version__} | " f"rasterio {rasterio.__version__} | geopandas {gpd.__version__} | overturemaps {version('overturemaps')} | " f"numpy {np.__version__}") def section(title: str) -> None: print(f"\n=== {title} ===") # --- Step 1: the area of interest and a chip-aligned 10 m grid ---------------------------------- NAME = "nottingham" CENTRE = (-1.15, 52.955) # lon/lat: Nottingham city centre CHIP = 128 # chip size in pixels: 1.28 km at 10 m NCHIPS = 24 # 24 x 24 chips = 30.72 x 30.72 km RES, CRS = 10, "EPSG:32630" # Sentinel-2's native 10 m grid in UTM zone 30N YEAR = 2021 # the year of ESA WorldCover v200, our land cover labels cx, cy = Transformer.from_crs(4326, CRS, always_xy=True).transform(*CENTRE) size = CHIP * NCHIPS * RES x0, y1 = round(cx - size / 2, -1), round(cy + size / 2, -1) # snap to whole 10 m geobox = GeoBox.from_bbox((x0, y1 - size, x0 + size, y1), crs=CRS, resolution=RES) TRANSFORM = from_origin(x0, y1, RES, RES) H = W = CHIP * NCHIPS bbox_ll = tuple(gpd.GeoSeries([box(x0, y1 - size, x0 + size, y1)], crs=CRS).to_crs(4326).total_bounds) section("area of interest") print(f"{NAME}: {H} x {W} pixels of {RES} m ({size / 1000:.2f} km square) in {CRS}; " f"lon/lat bounds {tuple(round(float(v), 3) for v in bbox_ll)}") # --- Step 2: a cloud-free Sentinel-2 L2A summer composite --------------------------------------- STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1" BANDS = ["B01", "B02", "B03", "B04", "B05", "B06", "B07", "B08", "B8A", "B09", "B11", "B12"] # L2A has no B10 catalog = pystac_client.Client.open(STAC_URL, modifier=planetary_computer.sign_inplace) s2_tif = DATA / f"s2_{NAME}_{YEAR}.tif" if not s2_tif.exists(): items = catalog.search(collections=["sentinel-2-l2a"], bbox=bbox_ll, datetime=f"{YEAR}-05-01/{YEAR}-08-31", query={"eo:cloud_cover": {"lt": 40}}).item_collection() section("Sentinel-2 search") print(f"{len(items)} L2A items with < 40% cloud, May to August {YEAR}, tiles " f"{sorted({i.properties['s2:mgrs_tile'] for i in items})}") odc.stac.configure_rio(cloud_defaults=True) ds = odc.stac.load(items, bands=BANDS + ["SCL"], geobox=geobox, groupby="solar_day", dtype="uint16", nodata=0, chunks={"time": 1, "x": 1536, "y": 1536}, resampling={"*": "bilinear", "SCL": "nearest"}) # Processing baseline 04.00 (from January 2022) adds 1000 to every L2A value; none in 2021, but check. base = {i.datetime.date(): float(i.properties["s2:processing_baseline"]) for i in items} offset = xr.DataArray([1000 if base[pd.Timestamp(t).date()] >= 4.0 else 0 for t in ds.time.values], dims="time", coords={"time": ds.time}) t0 = time.time() scl = ds["SCL"].compute() cloudy = np.isin(scl.values, [3, 8, 9, 10]) # cloud shadow, cloud, cirrus cloudy = ndimage.binary_dilation(cloudy, iterations=3, axes=(1, 2)) # grow by 30 m clear = scl.isin([4, 5, 6, 7]) & ~xr.DataArray(cloudy, coords=scl.coords, dims=scl.dims) frac = clear.mean(("y", "x")).to_series() keep = frac.index[frac >= 0.2] print(f"{len(frac)} acquisition dates; {len(keep)} are at least 20% clear: " f"{', '.join(f'{t:%d %b}' for t in keep)}") refl = (ds[BANDS].sel(time=keep).astype("float32") - offset.sel(time=keep)).where(clear.sel(time=keep)) med = refl.chunk({"time": -1, "y": 512, "x": 512}).median("time").compute() n_clear = clear.sel(time=keep).sum("time").values # SCL calls some bright roofs "cloud" on every date. Fill those gaps with the unmasked median. raw = (ds[BANDS].sel(time=keep).astype("float32") - offset.sel(time=keep)).where(ds["SCL"].sel(time=keep) > 0) med = med.fillna(raw.chunk({"time": -1, "y": 512, "x": 512}).median("time").compute()) print(f"pixels never clear, filled from the unmasked median: {(n_clear == 0).mean():.2%}") stack = np.stack([med[b].values for b in BANDS]) stack = np.where(np.isnan(stack), 0, np.clip(stack, 1, 65535)).astype("uint16") # 0 = no clear view with rasterio.open(s2_tif, "w", driver="GTiff", height=H, width=W, count=len(BANDS), dtype="uint16", crs=CRS, transform=TRANSFORM, nodata=0, tiled=True, compress="deflate") as dst: dst.write(stack) dst.descriptions = tuple(BANDS) print(f"median composite written in {time.time() - t0:.0f} s; clear views per pixel: " f"min {n_clear.min()}, median {np.median(n_clear):.0f}, max {n_clear.max()}") with rasterio.open(s2_tif) as src: s2 = src.read() # (12, H, W) surface reflectance x 10000 no_data = (s2 == 0).all(0) section("composite") print(f"{s2_tif.name}: {s2.shape[0]} bands, {s2.nbytes / 1e6:.0f} MB in memory, " f"pixels with no data: {no_data.mean():.3%}") # --- Step 3: label source 1, ESA WorldCover 2021 -------------------------------------------------- wc_items = catalog.search(collections=["esa-worldcover"], bbox=bbox_ll, query={"esa_worldcover:product_version": {"eq": "2.0.0"}}).item_collection() wc = odc.stac.load(wc_items, bands=["map"], geobox=geobox, resampling="nearest")["map"].isel(time=0).values section("ESA WorldCover 2021 (v200)") names = {10: "Tree cover", 20: "Shrubland", 30: "Grassland", 40: "Cropland", 50: "Built-up", 60: "Bare", 80: "Water", 90: "Wetland", 100: "Moss"} codes, counts = np.unique(wc, return_counts=True) print(" | ".join(f"{names.get(int(c), c)} {n / wc.size:.1%}" for c, n in zip(codes, counts))) # --- Step 4: label source 2, Overture Maps building footprints ---------------------------------- bld_file = DATA / f"overture_buildings_{NAME}.parquet" if not bld_file.exists(): t0 = time.time() b = overture.geodataframe("building", bbox=bbox_ll, stac=True).set_crs(4326, allow_override=True) b[["id", "geometry"]].assign(source=b["sources"].map(lambda s: s[0]["dataset"])).to_parquet(bld_file) print(f"downloaded {len(b):,} footprints in {time.time() - t0:.0f} s") bld = gpd.read_parquet(bld_file).to_crs(CRS) section(f"Overture Maps buildings (latest release {overture.get_latest_release()})") print(f"{len(bld):,} footprints; by source: " + ", ".join(f"{k} {v:.1%}" for k, v in bld["source"].value_counts(normalize=True).items())) # Rasterise at 2 m, then average 5 x 5 blocks: the share of each 10 m pixel covered by buildings. SUB = 5 fine = rasterio.features.rasterize(((g, 1) for g in bld.geometry), out_shape=(H * SUB, W * SUB), transform=from_origin(x0, y1, RES / SUB, RES / SUB), dtype="uint8") cover = fine.reshape(H, SUB, W, SUB).mean((1, 3)) del fine print(f"pixels touched by a building: {(cover > 0).mean():.1%}; at least half covered: {(cover >= 0.5).mean():.1%}") # --- Step 5: fuse the two sources into five classes ----------------------------------------------- CLASSES = ["Building", "Other built-up", "Trees", "Grass and crops", "Water"] COLOURS = ["#c21f3a", "#b0aaa0", "#1f5f3a", "#c9d98a", "#1543b0"] IGNORE = 255 label = np.full((H, W), IGNORE, "uint8") label[np.isin(wc, [20, 30, 40, 90, 100])] = 3 # shrubs, grass, crops, wetland, moss label[wc == 10] = 2 label[np.isin(wc, [50, 60])] = 1 # built-up and bare: roads, car parks, yards, quarries label[wc == 80] = 4 label[cover >= 0.5] = 0 # a building footprint wins wherever it covers half a pixel label[no_data] = IGNORE label_tif = DATA / f"labels_{NAME}_{YEAR}.tif" with rasterio.open(label_tif, "w", driver="GTiff", height=H, width=W, count=1, dtype="uint8", crs=CRS, transform=TRANSFORM, nodata=IGNORE, tiled=True, compress="deflate") as dst: dst.write(label, 1) dst.write_colormap(1, {i: tuple(int(c[k:k + 2], 16) for k in (1, 3, 5)) + (255,) for i, c in enumerate(COLOURS)}) section("fused label raster") print(pd.Series({c: (label == i).mean() for i, c in enumerate(CLASSES)} | {"ignored": (label == IGNORE).mean()}) .map("{:.1%}".format).to_string()) # --- Step 6: cut georeferenced chips ---------------------------------------------------------- rows = [] for r in range(NCHIPS): for c in range(NCHIPS): win = np.s_[r * CHIP:(r + 1) * CHIP, c * CHIP:(c + 1) * CHIP] lab = label[win] cx0, cy1 = x0 + c * CHIP * RES, y1 - r * CHIP * RES rec = {"chip_id": f"{NAME}_r{r:02d}_c{c:02d}", "row": r, "col": c, "valid": (lab != IGNORE).mean(), **{CLASSES[k]: (lab == k).mean() for k in range(len(CLASSES))}, "geometry": box(cx0, cy1 - CHIP * RES, cx0 + CHIP * RES, cy1)} rows.append(rec) if rec["valid"] < 0.95: continue profile = dict(driver="GTiff", height=CHIP, width=CHIP, crs=CRS, transform=from_origin(cx0, cy1, RES, RES), compress="deflate") with rasterio.open(CHIPS / "images" / f"{rec['chip_id']}.tif", "w", count=len(BANDS), dtype="uint16", nodata=0, **profile) as dst: dst.write(s2[(slice(None),) + win]) dst.descriptions = tuple(BANDS) with rasterio.open(CHIPS / "masks" / f"{rec['chip_id']}.tif", "w", count=1, dtype="uint8", nodata=IGNORE, **profile) as dst: dst.write(lab, 1) chips = gpd.GeoDataFrame(rows, crs=CRS) dropped = int((chips["valid"] < 0.95).sum()) chips = chips[chips["valid"] >= 0.95].reset_index(drop=True) section("chips") print(f"{len(chips)} chips of {CHIP} x {CHIP} pixels written; {dropped} dropped for > 5% missing data") # --- Step 7: split by spatial block, not at random -------------------------------------------- BLOCK = 4 # 4 x 4 chips = 5.12 km blocks chips["block"] = (chips["row"] // BLOCK) * (NCHIPS // BLOCK) + chips["col"] // BLOCK blocks = np.sort(chips["block"].unique()) area = chips[CLASSES].mean() def neighbour_leak(split: pd.Series) -> float: """Share of test chips that share an edge or corner with a training chip.""" train = set(zip(chips.loc[split == "train", "row"], chips.loc[split == "train", "col"])) test = chips[split == "test"] return float(np.mean([any((r + dr, c + dc) in train for dr in (-1, 0, 1) for dc in (-1, 0, 1)) for r, c in zip(test["row"], test["col"])])) # Try many random block assignments (70/15/15) and keep the one whose class shares best match the whole area. best = None n_hold = round(0.15 * len(blocks)) for seed in range(500): order = np.random.default_rng(seed).permutation(blocks) assign = {b: "val" for b in order[:n_hold]} | {b: "test" for b in order[n_hold:2 * n_hold]} split = chips["block"].map(lambda b: assign.get(b, "train")) shares = chips.groupby(split)[CLASSES].mean() score = float((np.abs(shares - area) / area).max().max()) # worst relative deviation of any class if best is None or score < best[0]: best = (score, seed, split) chips["split"] = best[2] random_split = pd.Series(np.random.default_rng(0).choice(["train", "val", "test"], len(chips), p=[0.7, 0.15, 0.15]), index=chips.index) leak_random, leak_block = neighbour_leak(random_split), neighbour_leak(chips["split"]) section("spatial split") print(f"{len(blocks)} blocks of {BLOCK} x {BLOCK} chips; best of 500 block shuffles: seed {best[1]} " f"(worst class share off by {best[0]:.0%})") print(chips["split"].value_counts().reindex(["train", "val", "test"]).to_string()) print(f"test chips touching a training chip: random split {leak_random:.0%}, block split {leak_block:.0%}") held = set(zip(chips.loc[chips["split"] != "train", "row"], chips.loc[chips["split"] != "train", "col"])) edge = pd.Series([any((r + dr, c + dc) in held for dr in (-1, 0, 1) for dc in (-1, 0, 1)) for r, c in zip(chips["row"], chips["col"])], index=chips.index) & (chips["split"] == "train") print(f"a one-chip buffer would remove {edge.sum()} of {(chips['split'] == 'train').sum()} training chips") # --- Step 8: class balance ----------------------------------------------------------------------- balance = chips.groupby("split")[CLASSES].mean().reindex(["train", "val", "test"]) balance.loc["all chips"] = chips[CLASSES].mean() shares = balance.div(balance.sum(1), axis=0) * 100 section("class balance (% of labelled pixels)") print(shares.round(1).to_string()) counts = np.bincount(label[label != IGNORE].ravel(), minlength=len(CLASSES)) weights = 1 / np.sqrt(counts / counts.sum()) print("loss weights, 1 / sqrt(frequency) scaled to mean 1: " + ", ".join(f"{c} {w:.2f}" for c, w in zip(CLASSES, weights / weights.mean()))) index_file = CHIPS / "chips.gpkg" chips.drop(columns=["valid"]).to_file(index_file, driver="GPKG", layer="chips") print(f"\nsaved {index_file.relative_to(HERE).as_posix()} with columns " f"{list(chips.drop(columns=['valid', 'geometry']).columns)}") # --- Step 9: the same rasters as a TorchGeo dataset (optional) ----------------------------------- try: from torch.utils.data import DataLoader from torchgeo.datasets import RasterDataset, stack_samples from torchgeo.samplers import RandomGeoSampler except ImportError: print("\nTorchGeo is not installed; skipping the optional TorchGeo check (pip install torchgeo).") else: class Sentinel2Composite(RasterDataset): filename_glob = f"s2_{NAME}_*.tif" is_image = True class FusedLabels(RasterDataset): filename_glob = f"labels_{NAME}_*.tif" is_image = False dataset = Sentinel2Composite(paths=DATA) & FusedLabels(paths=DATA) sampler = RandomGeoSampler(dataset, size=CHIP, length=64) batch = next(iter(DataLoader(dataset, sampler=sampler, batch_size=8, collate_fn=stack_samples))) section(f"TorchGeo {version('torchgeo')}") print(f"RandomGeoSampler batch: image {tuple(batch['image'].shape)} {batch['image'].dtype}, " f"mask {tuple(batch['mask'].shape)} {batch['mask'].dtype}") # --- Figures ------------------------------------------------------------------------------------ cmap = ListedColormap(COLOURS) legend = [Patch(color=c, label=n) for c, n in zip(COLOURS, CLASSES)] def rgb(img: np.ndarray) -> np.ndarray: """True colour (B04, B03, B02) with a fixed 0-0.25 reflectance stretch and gentle gamma.""" out = np.clip(img[[3, 2, 1]].transpose(1, 2, 0) / 2500, 0, 1) ** 0.8 out[(img == 0).all(0)] = 1 return out extent = (x0 / 1000, (x0 + size) / 1000, (y1 - size) / 1000, y1 / 1000) def frame(ax, title): ax.set_title(title, loc="left", fontsize=9) ax.set_xlabel("UTM 30N easting (km)") ax.set_ylabel("northing (km)") fig, axes = plt.subplots(1, 2, figsize=(11, 5.2), constrained_layout=True) axes[0].imshow(rgb(s2), extent=extent) frame(axes[0], f"Sentinel-2 L2A median composite, May to August {YEAR}") axes[1].imshow(np.ma.masked_equal(label, IGNORE), cmap=cmap, vmin=-0.5, vmax=4.5, extent=extent, interpolation="nearest") frame(axes[1], "Labels: Overture buildings over ESA WorldCover 2021") axes[1].legend(handles=legend, loc="upper left", bbox_to_anchor=(1.01, 1), fontsize=8, frameon=False) fig.savefig(OUT / "01-composite-and-labels.png") plt.close(fig) # Six example chips: the one richest in each class, plus the most mixed chip. props = chips[CLASSES].to_numpy() entropy = -(props * np.log(np.clip(props, 1e-9, 1))).sum(1) picks = list(dict.fromkeys([chips.loc[chips[c].idxmax(), "chip_id"] for c in CLASSES] + [chips.loc[int(entropy.argmax()), "chip_id"]])) fig, axes = plt.subplots(2, len(picks), figsize=(2 * len(picks), 4.6), constrained_layout=True) for j, cid in enumerate(picks): with rasterio.open(CHIPS / "images" / f"{cid}.tif") as s, rasterio.open(CHIPS / "masks" / f"{cid}.tif") as m: img, msk = s.read(), m.read(1) axes[0, j].imshow(rgb(img)) axes[1, j].imshow(rgb(img)) axes[1, j].imshow(np.ma.masked_equal(msk, IGNORE), cmap=cmap, vmin=-0.5, vmax=4.5, alpha=0.6, interpolation="nearest") axes[0, j].set_title(f"{cid.split('_', 1)[1]} ({chips.set_index('chip_id').loc[cid, 'split']})", fontsize=8) for ax in axes[:, j]: ax.set_xticks([]) ax.set_yticks([]) axes[0, 0].set_ylabel("image") axes[1, 0].set_ylabel("image + labels") fig.legend(handles=legend, loc="lower center", ncol=5, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.07)) fig.savefig(OUT / "02-chips-with-masks.png") plt.close(fig) split_cols = {"train": "#9db36b", "val": "#e0a030", "test": "#1543b0"} fig, axes = plt.subplots(1, 2, figsize=(11, 5.2), constrained_layout=True) for ax, sp, title in [(axes[0], random_split, f"Random split: {leak_random:.0%} of test chips touch a training chip"), (axes[1], chips["split"], f"Block split: {leak_block:.0%} of test chips touch a training chip")]: ax.imshow(rgb(s2), extent=extent) for g, s in zip(chips.geometry, sp): gx0, gy0, gx1, gy1 = (v / 1000 for v in g.bounds) ax.add_patch(Rectangle((gx0, gy0), gx1 - gx0, gy1 - gy0, facecolor=split_cols[s], alpha=0.7, edgecolor="white", linewidth=0.4)) frame(ax, title) axes[1].legend(handles=[Patch(color=c, label=k) for k, c in split_cols.items()], loc="upper left", bbox_to_anchor=(1.01, 1), fontsize=8, frameon=False) fig.savefig(OUT / "03-random-vs-block-split.png") plt.close(fig) fig, ax = plt.subplots(figsize=(8, 3.6), constrained_layout=True) xs = np.arange(len(CLASSES)) n_chips = chips["split"].value_counts().to_dict() | {"all chips": len(chips)} for i, (name, col) in enumerate(zip(shares.index, ["#9db36b", "#e0a030", "#1543b0", "#4a5553"])): ax.bar(xs + (i - 1.5) * 0.2, shares.loc[name], width=0.2, color=col, label=f"{name} ({n_chips[name]} chips)") ax.set_xticks(xs, CLASSES) ax.set_ylabel("share of labelled pixels (%)") ax.set_yscale("log") ax.set_ylim(0.3, 100) ax.set_title("Class balance by split (log scale)", loc="left", fontsize=9) ax.legend(fontsize=8, frameon=False) ax.grid(axis="y", alpha=0.3) fig.savefig(OUT / "04-class-balance.png") plt.close(fig) print(f"\nfigures written to {OUT.name}/")