Train a Semantic Segmentation Model for Water Detection¶

This notebook demonstrates how to train semantic segmentation models for water detection.

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

In [ ]:

url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/waterbody-dataset.zip"

url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/waterbody-dataset.zip"

In [ ]:

out_folder = geoai.download_file(url)

out_folder = geoai.download_file(url)

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"{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,
    num_classes=2,  # background and water
    batch_size=8,
    num_epochs=3,
    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"{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, num_classes=2, # background and water batch_size=8, num_epochs=3, 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"{out_folder}/unet_models/training_history.pth",
    figsize=(15, 5),
    verbose=True,
)

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

Run inference¶

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.

In [ ]:

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"

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"

In [ ]:

# Run semantic segmentation inference
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=4,
)

# Run semantic segmentation inference 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=4, )

Visualize the results¶

In [ ]:

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

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