Source code: Training data for deep learning on satellite imagery: Sentinel-2 image chips in Python

"""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}/")