Land Use / Land Cover Segmentation Using Sentinel-2 and Random Forest

!mamba install --channel rapidsai --quiet --yes cuml

!pip install gdown

import glob
import pickle
import dask.dataframe as dd
import pandas as pd
import geopandas as gpd
import hvplot.xarray
import matplotlib.pyplot as plt
import numpy as np
import odc.stac
import rasterio.features
import rioxarray as rxr
import xarray as xr
from cuml import RandomForestClassifier
from dask import compute, delayed
from dask_ml.model_selection import train_test_split
from geocube.api.core import make_geocube
from pystac_client import Client
from shapely.geometry import Polygon, box, mapping, shape
from sklearn.metrics import (
    ConfusionMatrixDisplay,
    accuracy_score,
    classification_report,
)
from tqdm import tqdm

# THIS IS A LINK TO DOWNLOAD ROIs_v9.zip ONLY
!gdown "https://drive.google.com/uc?id=1EOBFi26CsyOp-NsZRz6kFsbqy98KugME"

# Download the administrative boundaries (2016_phc_vut_pid_4326.geojson)
!wget https://pacificdata.org/data/dataset/9dba1377-740c-429e-92ce-6a484657b4d9/resource/3d490d87-99c0-47fd-98bd-211adaf44f71/download/2016_phc_vut_pid_4326.geojson

ROI_version = 9

lulc_gdf = gpd.read_file(f"ROIs_v{ROI_version}.zip")

lulc_gdf

admin_boundaries_gdf = gpd.read_file("./2016_phc_vut_pid_4326.geojson")

admin_boundaries_gdf

len(lulc_gdf), len(admin_boundaries_gdf)

lulc_gdf.columns

lulc_gdf.ROI.unique()

lulc_gdf.sample(frac=1).head(10)

# original projection in units of degrees
lulc_gdf.crs

# convert to a projection with units in meters that is appropriate for Vanuatu
lulc_gdf_meters = lulc_gdf.to_crs(epsg="32759") # WGS 84 / UTM zone 59S

lulc_gdf_meters.crs

lulc_gdf_meters.dissolve(by=["ROI", "Pname"]).area.to_csv("Vanuatu_ROIs.csv")

lulc_gdf_meters.dissolve(by=["ROI", "Pname"]).area

YEAR = 2020 # year matching label data 
PROVINCES_TRAIN = ["TORBA", "SANMA", "PENAMA", "MALAMPA", "TAFEA", "SHEFA"] # all provinces

admin_boundaries_gdf = admin_boundaries_gdf.set_index(keys="pname")  # set province name as the index

GEOMS_TRAIN = admin_boundaries_gdf.loc[PROVINCES_TRAIN].geometry.tolist()
GEOMS_TRAIN[0]

STAC_URL = "http://stac.digitalearthpacific.org/"
stac_client = Client.open(STAC_URL)

# Collect s2_data per train province in a list
s2_data_train_list = []

for pname, geom in tqdm(zip(PROVINCES_TRAIN, GEOMS_TRAIN), total=len(GEOMS_TRAIN), desc="Loading GeoMAD per province"):
    try:
        # Query STAC for this province
        s2_search = stac_client.search(
            collections=["dep_s2_geomad"], # Sentinel-2 Geometric Median and Absolute Deviations (GeoMAD) over the Pacific.
            intersects=mapping(geom),  # GeoJSON dict
            datetime=str(YEAR),
        )
        s2_items = s2_search.item_collection()

        if len(s2_items) == 0:
            print(f"No items found for {pname}")
            continue

        # Load data from items
        s2_data = odc.stac.load(
            items=s2_items,
            bands=["blue", "green", "red", "nir08", "swir16"],
            chunks={"x": 1024, "y": 1024, "bands": -1, "time": -1},
            crs="EPSG:32759",
            resolution=20,
        )

        s2_data_train_list.append(s2_data)

    except Exception as e:
        print(f"Error loading {pname}: {e}")

# Should be the same projection as use dearlier (WGS 84 / UTM zone 59S) with units in meters
s2_data_train_list[0].rio.crs

lulc_gdf_meters.Pname.unique()

# Create a dictionary of buffered geometries per province
admin_boundaries_gdf_meters = admin_boundaries_gdf.to_crs(crs=s2_data_train_list[0].rio.crs)

geom_buffered_train = {
    pname: admin_boundaries_gdf_meters.loc[pname].geometry.buffer(5) # minimal; 5 meters only
    for pname in PROVINCES_TRAIN
}

geom_train_buffered_list = list(geom_buffered_train.values())
geom_train_buffered_list

geom_train_buffered_list[0]

# Make sure the keys match — we'll zip province names, geometries, and s2 datasets
s2_train_clipped_list = []

for pname, geom, s2_data in zip(PROVINCES_TRAIN, geom_buffered_train.values(), s2_data_train_list):
    try:
        # Clip the dataset to the buffered province geometry
        s2_clipped = s2_data.rio.clip(
            geometries=[mapping(geom)],
            crs=s2_data.rio.crs,
            drop=True
        )
        s2_train_clipped_list.append(s2_clipped)
    except Exception as e:
        print(f"Error clipping data for {pname}: {e}")

s2_train_clipped_list[0]

# Plot sample train province
s2_rgb = s2_train_clipped_list[0][["red", "green", "blue"]] 
s2_rgb_array = s2_rgb.to_array("band")  # now dims: band, y, x
s2_rgb_array_squeezed = s2_rgb_array.squeeze(dim="time", drop=True)

s2_rgb_array_squeezed.plot.imshow(size=4, vmin=0, vmax=4000)

# Calculate remote sensing indices useful for mapping LULC
def compute_indices(ds):
    red = ds["red"]
    green = ds["green"]
    blue = ds["blue"]
    nir = ds["nir08"]
    swir = ds["swir16"]
    eps = 1e-6
    return xr.Dataset({
        "NDVI": (nir - red) / (nir + red + eps),
        "MNDWI": (green - swir) / (green + swir + eps),
        "SAVI": ((nir - red) / (nir + red + 0.5 + eps)) * 1.5,
        "BSI": ((swir + red) - (nir + blue)) / ((swir + red) + (nir + blue) + eps),
        "NDBI": (swir - nir) / (swir + nir + eps),
    })

index_data_train_list = []

for s2_clipped in s2_train_clipped_list:
    index_data = compute_indices(s2_clipped).squeeze("time", drop=True)
    index_data_train_list.append(index_data)

# Rasterize labels
rasterized_labels_train_dict = {}
gdf_train_clipped_dict = {}
metadata_dict = {}

for s2_clipped_train, pname in zip(s2_train_clipped_list, PROVINCES_TRAIN):
    print(f"\n Processing province: {pname}")

    # Sentinel-2 metadata
    width = s2_clipped_train.sizes['x']
    height = s2_clipped_train.sizes['y']
    resolution = s2_clipped_train.rio.resolution()
    bounds = s2_clipped_train.rio.bounds()
    epsg = s2_clipped_train.rio.crs.to_epsg()
    raster_bounds = box(*bounds)

    # Class mapping (safe per-province if needed)
    unique_classes = lulc_gdf_meters['ROI'].unique()
    #class_mapping = {cls: i+1 for i, cls in enumerate(unique_classes)}
    class_mapping = {cls: i for i, cls in enumerate(unique_classes)} # zero-based, assumes existence of no data    
    lulc_gdf_meters['ROI_numeric'] = lulc_gdf_meters['ROI'].map(class_mapping)

    # Clip vector LULC to S2 bounds
    gdf_train_clipped = lulc_gdf_meters[lulc_gdf_meters.intersects(raster_bounds)]
    print(f"Vector features: {len(lulc_gdf_meters)} → Clipped: {len(gdf_train_clipped)}")

    if len(gdf_train_clipped) == 0:
        print(f"No vector data found for province: {pname}, skipping rasterization.")
        continue

    # Rasterize clipped vector labels
    rasterized_labels_train = make_geocube(
        vector_data=gdf_train_clipped,
        measurements=["ROI_numeric"],
        like=s2_clipped_train
    )

    # Store outputs
    rasterized_labels_train_dict[pname] = rasterized_labels_train
    gdf_train_clipped_dict[pname] = gdf_train_clipped
    metadata_dict[pname] = {
        "width": width,
        "height": height,
        "epsg": epsg,
        "resolution": resolution,
        "bounds": bounds,
        "class_mapping": class_mapping
    }

fig, ax = plt.subplots()
gdf_train_clipped_dict["TORBA"].plot(ax=ax, facecolor="none", edgecolor="red")
bbox = box(*s2_train_clipped_list[0].rio.bounds())
gpd.GeoSeries([bbox], crs=s2_train_clipped_list[0].rio.crs).plot(ax=ax, facecolor="none", edgecolor="blue")
plt.title("gdf_train (red) vs. Raster Bounds (blue)")
plt.show()

gdf_train_clipped_dict["TORBA"].ROI.unique()

features_list = []
labels_list = []

for i, prov_name in enumerate(PROVINCES_TRAIN):
    # Stack spatial dimensions first
    features_train = index_data_train_list[i].to_array().stack(flattened_pixel=("y", "x"))
    labels_train = rasterized_labels_train_dict[prov_name].to_array().stack(flattened_pixel=("y", "x"))
    
    # Compute mask for valid pixels (no NaNs across all features or labels)
    mask = (
        np.isfinite(features_train).all(dim="variable") &
        np.isfinite(labels_train).all(dim="variable")
    ).compute()
    
    # Apply the mask to drop invalid pixels
    features_train = features_train[:, mask].transpose("flattened_pixel", "variable").compute()
    labels_train = labels_train[:, mask].transpose("flattened_pixel", "variable").squeeze().compute()
    
    labels_train = labels_train.astype(int)

    features_list.append(features_train)
    labels_list.append(labels_train)

# Concatenate all provinces along the flattened_pixel dimension
features_train = xr.concat(features_list, dim="flattened_pixel")
labels_train = xr.concat(labels_list, dim="flattened_pixel")

print("Combined features_train shape:", features_train.shape)
print("Combined labels_train shape:", labels_train.shape)

np.unique(labels_train.values)

len(features_train), len(labels_train)

X_train, X_test, y_train, y_test = train_test_split(
    features_train, labels_train, test_size=0.2, random_state=42, shuffle=True
)

np.unique(y_train), np.unique(y_test) 

len(X_train), len(X_test), len(y_train), len(y_test)

%%time
# Train a Random Forest classifier
clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train.data, y_train.data)

# Test the classifier
y_pred = clf.predict(X_test)

accuracy = accuracy_score(y_test, y_pred)
print(f"Accuracy: {accuracy:.4f}")

print(classification_report(y_true=y_test, y_pred=y_pred))

print("Class mapping:")
for key, val in class_mapping.items():
    print(val, key)

ConfusionMatrixDisplay.from_predictions(
    y_true=y_test, y_pred=y_pred, normalize="true", values_format=".2f"
)

pickle.dump(clf, open(f"rf_vanuatu_lulc_{YEAR}_ROIversion{ROI_version}.pkl", "wb"))