Water Detection with Sentinel-2 Imagery¶

This notebook demonstrates how to train semantic segmentation models for water detection using Sentinel-2 imagery.

Install packages¶

To use the new functionality, ensure the required packages are installed.

In [ ]:

# %pip install geoai-py

# %pip install geoai-py

Import libraries¶

In [ ]:

import geoai

import geoai

Download sample data¶

We'll use the Earth Surface Water Dataset from Zenodo. Credits to the author (Xin Luo) of the dataset

In [ ]:

url = "https://zenodo.org/records/5205674/files/dset-s2.zip?download=1"

url = "https://zenodo.org/records/5205674/files/dset-s2.zip?download=1"

In [ ]:

data_dir = geoai.download_file(url)

data_dir = geoai.download_file(url)

In [ ]:

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"

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.

In [ ]:

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,
)

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. 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:

In [ ]:

# Test train_segmentation_model with automatic size detection
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=8,
    num_epochs=50,
    learning_rate=0.001,
    val_split=0.2,
    verbose=True,
)

# Test train_segmentation_model with automatic size detection 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=8, num_epochs=50, learning_rate=0.001, val_split=0.2, verbose=True, )

Evaluate the model¶

Let's examine the training curves and model performance:

In [ ]:

geoai.plot_performance_metrics(
    history_path=f"{tiles_dir}/unet_models/training_history.pth",
    figsize=(15, 5),
    verbose=True,
)

geoai.plot_performance_metrics( history_path=f"{tiles_dir}/unet_models/training_history.pth", figsize=(15, 5), verbose=True, )

Run inference¶

In [ ]:

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"

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"

In [ ]:

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=8,
    quiet=True,
)

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=8, quiet=True, )

Visualize results¶

In [ ]:

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,
)

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, )