Automating surface water mapping with AI tools¶
This notebook is developed for the Australian Water School premium webinar on automating surface water mapping with AI tools on July 23, 2025. To register for the webinar, please visit https://awschool.com.au/training/ai-tools-for-mapping.
More resources:
- GitHub: https://github.com/opengeos/geoai
- Documentation: https://opengeoai.org
- Notebook: https://opengeoai.org/workshops/GeoAI_Workshop_2025
- Web App: https://huggingface.co/spaces/giswqs/surface-water-app
- Web App Demo: https://youtu.be/LKySb6OYU7M
- YouTube Tutorials: https://tinyurl.com/GeoAI-Tutorials
Use Colab GPU¶
To use GPU, please click the "Runtime" menu and select "Change runtime type". Then select "T4 GPU" from the dropdown menu.
Install packages¶
Uncomment the following cell to install the package. It may take a few minutes to install the package. Please be patient.
# %pip install geoai-py
Import libraries¶
Next, we need to install the geoai
package.
import geoai
Surface water mapping with non-georeferenced satellite imagery¶
In the first part of this notebook, we will demonstrate how to map surface water using non-georeferenced satellite imagery in jpg/png format.
Download sample data¶
We'll use the waterbody dataset from Kaggle. You will need to create an account and download the dataset. I have already downloaded the dataset and saved a copy on Hugging Face. Let's download the dataset:
url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/waterbody-dataset.zip"
out_folder = geoai.download_file(url)
The unzipped dataset contains two folders: images
and masks
. Each folder contains 2,841 images in jpg format. The images
folder contains the original satellite imagery, and the masks
folder contains the corresponding surface water masks. Note that the image size varies greatly.
We will use these 2,841 image pairs to train a surface water mapping model.
Train semantic segmentation model¶
Before diving into the training, let's understand the basic concepts of deep learning architectures and encoders.
Deep learning architecture is like the blueprint for a neural network — it defines how the network is built and how data flows through it. It includes layers of nodes (neurons) that process the data step by step to learn patterns, like identifying cats in pictures or translating languages.
There are many types of architectures:
- Feedforward Neural Networks (simple, goes one way)
- Convolutional Neural Networks (CNNs) (used for images)
- Recurrent Neural Networks (RNNs) (used for sequences like speech)
- Transformers (used for language tasks, like ChatGPT)
Each one is designed for specific types of problems.
An encoder is a part of the neural network that takes the input (like a sentence or image) and compresses it into a smaller, meaningful form called a feature representation or embedding. It captures the important information while throwing away the noise.
For example, if you feed the sentence “I love pizza” into the encoder, it turns that sentence into a set of numbers that still represent its meaning but are easier for the computer to understand and work with.
Encoders are used in models like:
- Autoencoders (for compressing and reconstructing data)
- Transformer Encoders (like BERT, for understanding language)
- Encoder-Decoder models (like translation systems)
Now we'll train a semantic segmentation model using the new train_segmentation_model
function. This function supports various architectures from segmentation-models-pytorch
:
- Architectures:
unet
,unetplusplus
deeplabv3
,deeplabv3plus
,fpn
,pspnet
,linknet
,manet
- Encoders:
resnet34
,resnet50
,efficientnet-b0
,mobilenet_v2
, etc.
For more details, please refer to the segmentation-models-pytorch documentation.
Let's train the module using U-Net with ResNet34 encoder:
# Test train_segmentation_model with automatic size detection
geoai.train_segmentation_model(
images_dir=f"{out_folder}/images",
labels_dir=f"{out_folder}/masks",
output_dir=f"{out_folder}/unet_models",
architecture="unet",
encoder_name="resnet34",
encoder_weights="imagenet",
num_channels=3, # number of channels in the input image
num_classes=2, # background and water
batch_size=32, # The number of images to process in each batch
num_epochs=3, # training for 3 epochs to save time, in practice you should train for more epochs
learning_rate=0.001, # learning rate for the optimizer
val_split=0.2, # 20% of the data for validation
target_size=(512, 512), # target size of the input image
verbose=True, # print progress
)
In the model output folder unet_models
, you will find the following files:
best_model.pth
: The best model checkpointfinal_model.pth
: The last model checkpointtraining_history.pth
: The training historytraining_summary.txt
: The training summary
Evaluate the model¶
Let's examine the training curves and model performance:
geoai.plot_performance_metrics(
history_path=f"{out_folder}/unet_models/training_history.pth",
figsize=(15, 5),
verbose=True,
)
Run inference on a single image¶
You can run inference on a new image using the semantic_segmentation
function. I don't have a new image to test on, so I'll use one of the training images. In reality, you would use your own images not used in training.
index = 3
test_image_path = f"{out_folder}/images/water_body_{index}.jpg"
ground_truth_path = f"{out_folder}/masks/water_body_{index}.jpg"
prediction_path = f"{out_folder}/prediction/water_body_{index}.png" # save as png to preserve exact values and avoid compression artifacts
model_path = f"{out_folder}/unet_models/best_model.pth"
geoai.semantic_segmentation(
input_path=test_image_path,
output_path=prediction_path,
model_path=model_path,
architecture="unet",
encoder_name="resnet34",
num_channels=3,
num_classes=2,
window_size=512,
overlap=256,
batch_size=32,
)
fig = geoai.plot_prediction_comparison(
original_image=test_image_path,
prediction_image=prediction_path,
ground_truth_image=ground_truth_path,
titles=["Original", "Prediction", "Ground Truth"],
figsize=(15, 5),
save_path=f"{out_folder}/prediction/water_body_{index}_comparison.png",
show_plot=True,
)
Run inference on multiple images¶
First, let's download the test images and masks.
url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/waterbody-dataset-sample.zip"
data_dir = geoai.download_file(url)
images_dir = f"{data_dir}/images"
masks_dir = f"{data_dir}/masks"
predictions_dir = f"{data_dir}/predictions"
geoai.semantic_segmentation_batch(
input_dir=images_dir,
output_dir=predictions_dir,
model_path=model_path,
architecture="unet",
encoder_name="resnet34",
num_channels=3,
num_classes=2,
window_size=512,
overlap=256,
batch_size=4,
quiet=True,
)
Surface water mapping with Sentinel-2 imagery¶
In the second part of this notebook, we will demonstrate how to map surface water using Sentinel-2 imagery with six spectral bands, including blue, green, red, near-infrared, and short-wave infrared bands.
Download sample data¶
We'll use the Earth Surface Water Dataset from Zenodo. Credits to the author (Xin Luo) of the dataset
url = "https://zenodo.org/records/5205674/files/dset-s2.zip?download=1"
data_dir = geoai.download_file(url)
In the unzipped dataset, we have four folders:
dset-s2/tra_scene
: training imagesdset-s2/tra_truth
: training masksdset-s2/val_scene
: validation imagesdset-s2/val_truth
: validation masks
We will use the training images and masks to train a semantic segmentation model.
images_dir = f"{data_dir}/dset-s2/tra_scene"
masks_dir = f"{data_dir}/dset-s2/tra_truth"
tiles_dir = f"{data_dir}/dset-s2/tiles"
Create training data¶
We'll create the same training tiles as before.
result = geoai.export_geotiff_tiles_batch(
images_folder=images_dir,
masks_folder=masks_dir,
output_folder=tiles_dir,
tile_size=512,
stride=128,
quiet=True,
)
Train semantic segmentation model¶
Now we'll train a semantic segmentation model using the new train_segmentation_model
function. Let's train the module using U-Net with ResNet34 encoder:
geoai.train_segmentation_model(
images_dir=f"{tiles_dir}/images",
labels_dir=f"{tiles_dir}/masks",
output_dir=f"{tiles_dir}/unet_models",
architecture="unet",
encoder_name="resnet34",
encoder_weights="imagenet",
num_channels=6,
num_classes=2, # background and water
batch_size=32,
num_epochs=5, # training for 5 epochs to save time, in practice you should train for more epochs
learning_rate=0.001,
val_split=0.2,
verbose=True,
)
Evaluate the model¶
Let's examine the training curves and model performance:
geoai.plot_performance_metrics(
history_path=f"{tiles_dir}/unet_models/training_history.pth",
figsize=(15, 5),
verbose=True,
)
Run inference¶
images_dir = f"{data_dir}/dset-s2/val_scene"
masks_dir = f"{data_dir}/dset-s2/val_truth"
predictions_dir = f"{data_dir}/dset-s2/predictions"
model_path = f"{tiles_dir}/unet_models/best_model.pth"
geoai.semantic_segmentation_batch(
input_dir=images_dir,
output_dir=predictions_dir,
model_path=model_path,
architecture="unet",
encoder_name="resnet34",
num_channels=6,
num_classes=2,
window_size=512,
overlap=256,
batch_size=32,
quiet=True,
)
Visualize results¶
test_image_path = (
f"{data_dir}/dset-s2/val_scene/S2A_L2A_20190318_N0211_R061_6Bands_S2.tif"
)
ground_truth_path = (
f"{data_dir}/dset-s2/val_truth/S2A_L2A_20190318_N0211_R061_S2_Truth.tif"
)
prediction_path = (
f"{data_dir}/dset-s2/predictions/S2A_L2A_20190318_N0211_R061_6Bands_S2_mask.tif"
)
save_path = f"{data_dir}/dset-s2/S2A_L2A_20190318_N0211_R061_6Bands_S2_comparison.png"
fig = geoai.plot_prediction_comparison(
original_image=test_image_path,
prediction_image=prediction_path,
ground_truth_image=ground_truth_path,
titles=["Original", "Prediction", "Ground Truth"],
figsize=(15, 5),
save_path=save_path,
show_plot=True,
indexes=[5, 4, 3],
divider=5000,
)
Download Sentinel-2 imagery¶
m = geoai.Map(center=[-16.3043, 128.7412], zoom=10)
m.add_basemap("Esri.WorldImagery")
m.add_stac_gui()
m
# m.stac_gdf
# m.stac_item
import leafmap
url = "https://earth-search.aws.element84.com/v1/"
collection = "sentinel-2-l2a"
time_range = "2025-01-01/2025-07-20"
bbox = [128.6735, -16.2466, 128.9577, -16.0962]
search = leafmap.stac_search(
url=url,
max_items=10,
collections=[collection],
bbox=bbox,
datetime=time_range,
query={"eo:cloud_cover": {"lt": 10}},
sortby=[{"field": "properties.eo:cloud_cover", "direction": "asc"}],
get_collection=True,
)
search
search = leafmap.stac_search(
url=url,
max_items=10,
collections=[collection],
bbox=bbox,
datetime=time_range,
query={"eo:cloud_cover": {"lt": 10}},
sortby=[{"field": "properties.eo:cloud_cover", "direction": "asc"}],
get_gdf=True,
)
search.head()
search = leafmap.stac_search(
url=url,
max_items=1,
collections=[collection],
bbox=bbox,
datetime=time_range,
query={"eo:cloud_cover": {"lt": 10}},
sortby=[{"field": "properties.eo:cloud_cover", "direction": "asc"}],
get_assets=True,
)
search
bands = ["blue", "green", "red", "nir", "swir16", "swir22"]
assets = list(search.values())[0]
links = [assets[band] for band in bands]
for link in links:
print(link)
out_dir = "s2"
leafmap.download_files(links, out_dir)
Stack image bands¶
Uncomment the following cell to install GDAL on Colab.
# !apt-get install -y gdal-bin
s2_path = "s2.tif"
geoai.stack_bands(input_files=out_dir, output_file=s2_path)
geoai.view_raster(s2_path, indexes=[4, 3, 2])
Run inference on a Sentinel-2 image¶
s2_mask = "s2_mask.tif"
model_path = f"{tiles_dir}/unet_models/best_model.pth"
geoai.semantic_segmentation(
input_path=s2_path,
output_path=s2_mask,
model_path=model_path,
architecture="unet",
encoder_name="resnet34",
num_channels=6,
num_classes=2,
window_size=512,
overlap=256,
batch_size=32,
)
Visualize the results¶
geoai.view_raster(
s2_mask, no_data=0, colormap="viridis", basemap=s2_path, backend="ipyleaflet"
)
train_raster_url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip/naip_water_train.tif"
train_masks_url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip/naip_water_masks.tif"
test_raster_url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip/naip_water_test.tif"
train_raster_path = geoai.download_file(train_raster_url)
train_masks_path = geoai.download_file(train_masks_url)
test_raster_path = geoai.download_file(test_raster_url)
geoai.print_raster_info(train_raster_path, show_preview=False)
Visualize sample data¶
geoai.view_raster(train_masks_url, nodata=0, basemap=train_raster_url)
geoai.view_raster(test_raster_url)
Create training data¶
out_folder = "naip"
tiles = geoai.export_geotiff_tiles(
in_raster=train_raster_path,
out_folder=out_folder,
in_class_data=train_masks_path,
tile_size=512,
stride=128,
buffer_radius=0,
)
Train segmentation model¶
geoai.train_segmentation_model(
images_dir=f"{out_folder}/images",
labels_dir=f"{out_folder}/labels",
output_dir=f"{out_folder}/models",
architecture="unet",
encoder_name="resnet34",
encoder_weights="imagenet",
num_channels=4,
pretrained=True,
batch_size=8,
num_epochs=5,
learning_rate=0.005,
val_split=0.2,
)
Evaluate the model¶
geoai.plot_performance_metrics(
history_path=f"{out_folder}/models/training_history.pth",
figsize=(15, 5),
verbose=True,
)
Run inference¶
masks_path = "naip_water_prediction.tif"
model_path = f"{out_folder}/models/best_model.pth"
geoai.semantic_segmentation(
test_raster_path,
masks_path,
model_path,
architecture="unet",
encoder_name="resnet34",
encoder_weights="imagenet",
window_size=512,
overlap=128,
confidence_threshold=0.3,
batch_size=32,
num_channels=4,
)
Vectorize masks¶
output_path = "naip_water_prediction.geojson"
gdf = geoai.raster_to_vector(
masks_path, output_path, min_area=1000, simplify_tolerance=1
)
gdf = geoai.add_geometric_properties(gdf)
len(gdf)
geoai.view_vector_interactive(gdf, tiles=test_raster_url)
gdf["elongation"].hist()
gdf_filtered = gdf[gdf["elongation"] < 10]
len(gdf_filtered)
Visualize results¶
geoai.view_vector_interactive(gdf_filtered, tiles=test_raster_url)
geoai.create_split_map(
left_layer=gdf_filtered,
right_layer=test_raster_url,
left_args={"style": {"color": "red", "fillOpacity": 0.2}},
basemap=test_raster_url,
)