Land Accounts inference on Sentinel-2 GeoMAD with Random Forest

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

import geopandas as gpd
import pickle
import hvplot.xarray
import matplotlib.pyplot as plt
import numpy as np
import odc.stac
import rasterio.features
from rasterio.features import sieve
import rioxarray
import xarray as xr
# from cuml import RandomForestClassifier
from pystac_client import Client
from shapely.geometry import box, mapping, shape

# Version of the LULC model (based on, for example, ROIs_v9.zip)
ROI_VERSION = "9"
YEAR = 2020  # year to run inference on
PROVINCE_INFERENCE = "SHEFA"  # Vanuatu province, choose from ["TORBA", "SANMA", "PENAMA", "MALAMPA", "TAFEA"]

# 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

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

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

admin_boundaries_gdf

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

admin_boundaries_gdf

GEOM_INFERENCE = admin_boundaries_gdf.loc[PROVINCE_INFERENCE].geometry
GEOM_INFERENCE

unique_classes = lulc_gdf["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
print(class_mapping)

gdf_inference = lulc_gdf.query(expr=f"Pname == '{PROVINCE_INFERENCE}'")

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

s2_search = stac_client.search(
    collections=["dep_s2_geomad"],
    intersects=GEOM_INFERENCE,
    datetime=str(YEAR),
)
# Retrieve all items from search results
s2_items = s2_search.item_collection()
print("len(s2_items): ", len(s2_items))

s2_data_inference = 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_inference

# Keep projection aligned with raster
raster_crs = s2_data_inference.rio.crs
print(raster_crs)

# Get only the select province and reproject
admin_boundaries_reprojected_inference = admin_boundaries_gdf.loc[[PROVINCE_INFERENCE]].to_crs(
    crs=raster_crs
)

# Buffer in raster units (meters if UTM)
geom_buffered_inference = admin_boundaries_reprojected_inference.buffer(distance=5)[PROVINCE_INFERENCE]
geom_buffered_inference

# Clip inference province
s2_clipped_inference = s2_data_inference.rio.clip(geometries=[geom_buffered_inference])

# Plot inference province
s2_rgb = s2_clipped_inference[["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=3000)

# 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_inference = compute_indices(s2_clipped_inference).squeeze("time", drop=True)
print(index_data_inference)

width_inference, height_inference = s2_clipped_inference.x.size, s2_clipped_inference.y.size
gdf_inference = gdf_inference.to_crs(epsg=s2_clipped_inference.rio.crs.to_epsg())

features_inference = index_data_inference.to_array().stack(flattened_pixel=("y", "x"))

features_inference = features_inference[:].transpose("flattened_pixel", "variable").compute()

print("features_inference shape:", features_inference.shape)

len(features_inference)

%%time
# Load trained Random Forest classifier
clf = pickle.load(open(f"rf_vanuatu_lulc_{YEAR}_ROIversion{ROI_VERSION}.pkl", "rb"))

y_pred = clf.predict(features_inference)

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

predicted_map = y_pred.reshape((height_inference, width_inference))
predicted_map.shape

predicted_map_xr = xr.DataArray(
    data=predicted_map,
    coords=s2_clipped_inference.squeeze("time", drop=True).coords,
)
print(np.unique(y_pred))

predicted_map_xr

predicted_map_xr.rio.crs

predicted_map_xr = predicted_map_xr.rio.clip(
    geometries=[geom_buffered_inference],
    crs=predicted_map_xr.rio.crs,
    drop=True
)

# Convert sentinel int value to NaN
predicted_map_xr = predicted_map_xr.where(predicted_map_xr != np.iinfo(predicted_map_xr.dtype).min)

np.unique(predicted_map_xr.values)

predicted_map_xr.hvplot.image(height=600, rasterize=True, cmap="Set1")

predicted_map_xr

def sieve_xarray_with_nodata(da, size=3):
    """
    Apply rasterio sieve to an xarray.DataArray with integer classes,
    remove small clumps, and set masked areas to NaN.
    """
    arr = da.values.copy()
    
    # Temporarily replace NaNs with a float value
    sentinel = -9999
    arr_filled = np.where(np.isnan(arr), sentinel, arr).astype(np.int32)
    
    # Apply sieve
    sieved = sieve(arr_filled, size=size, connectivity=8)
    
    # Convert to float before assigning NaN
    sieved = sieved.astype("float32")
    
    # Restore NaN for previously masked pixels
    sieved[sieved == sentinel] = np.nan
    
    return xr.DataArray(
        sieved,
        coords=da.coords,
        dims=da.dims,
        attrs=da.attrs
    )

lulc_smoothed = sieve_xarray_with_nodata(predicted_map_xr, size=3) # increase size to intensify smoothing

# Check unique values (ignoring NaNs)
print(np.unique(lulc_smoothed.values[~np.isnan(lulc_smoothed.values)]))

lulc_smoothed.hvplot.image(height=600, rasterize=True, cmap="Set1")

compatible_array = lulc_smoothed.astype("float32") #predicted_map_xr.astype("float32")
mask = ~np.isnan(compatible_array.values)  # mask valid values

# Rasterize to polygons
polygons = list(
    rasterio.features.shapes(
        compatible_array.values, transform=compatible_array.rio.transform(),
        mask=mask,
    )
)

# Convert polygons to GeoDataFrame and filter out masked value
prediction_gdf = gpd.GeoDataFrame(
    [{"geometry": shape(geom), "value": int(value)} for geom, value in polygons if not np.isnan(value)],
    crs=predicted_map_xr.rio.crs  # keep raster CRS
)

# print(prediction_gdf)
print(prediction_gdf.value.unique())

prediction_gdf.to_file(
    f"./predicted_lulc_utm_{PROVINCE_INFERENCE}_{YEAR}_ROIv{ROI_VERSION}_smoothed.geojson", driver="GeoJSON"
)

prediction_gdf.head(10)

# prediction_gdf_merged_nationwide = pd.concat(
#     [
#         prediction_gdf_TORBA,
#         prediction_gdf_SANMA,
#         prediction_gdf_PENAMA,
#         prediction_gdf_MALAMPA,
#         prediction_gdf_SHEFA,
#         prediction_gdf_TAFEA,
#     ],
#     ignore_index=True,
# )
#
# prediction_gdf_merged_nationwide.to_file(
#     f"./predicted_lulc_utm_nationwide_{YEAR}_ROIv{ROI_VERSION}.geojson", driver="GeoJSON"
# )