Detection of small water bodies for vector control using deep learning on multispectral imagery from unmanned aerial vehicles

This is the official github published datasets and implementation described in the following paper.

Detection of small water bodies for vector control using deep learning on multispectral imagery from unmanned aerial vehicles.

Phuc Linh Ngo, Viet Hoang Pham, Ngoc Long Bui, Huynh Anh Thu Phan, Bich Hien Vo, Thirumalaisamy P. Velavan, Duc Khanh Tran

WaterMAI Dataset Specification

Please download the published WaterMAI dataset from WaterMAI_dataset and orthomosaic images from WaterMAI Orthomosaic

• Folder co and ir in labels folder are the bounding box labels (YOLO type) of WaterMAI dataset. The masks of those 2 folders are the same, only the file's names are different
• Folder mask_co in labels folder is the segmentation label mask of WaterMAI dataset.

WaterMAI_dataset
├── images
│   ├── co
│   │   ├── val_co.txt
│   │   ├── train_co.txt
│   │   ├── DJI_1010_co.jpg
│   │   ├── ...
│   │   └── DJI_13970_co.jpg
│   └── ir
│       ├── val_ir.txt
│       ├── train_ir.txt
│       ├── DJI_1010_ir.jpg
│       ├── ...
│       └── DJI_13970_ir.jpg
└── labels
│    ├── color
│    │   ├── DJI_1010_co.txt
│    │   ├── ...
│    │   └── DJI_13970_co.txt
│    ├── ir
│    │   ├── DJI_1010_ir.txt
│    │   ├── ...
│    │   └── DJI_13970_ir.txt
│    └── mask_co
│        ├── DJI_1010_co.png
│        ├── ...
│        └── DJI_13970_co.png
│___testset
     ├── images
     │     ├── co
     │     │   ├── DJI_0010_102_co.jpg
     │     │   └── ...
     │     └── ir
     │         ├── DJI_0010_102_ir.jpg
     │         └── ...
     ├── labels
           ├── color
           │   ├── DJI_0010_102_co.txt
           │   └── ...
           ├── ir
           │   ├── DJI_0010_102_ir.txt
           │   └── ...
           └── mask_co
               ├── DJI_0010_102_co.png
               └── ...
           

Example

• 1st column contains raw RGB images
• 2nd column contains raw NIR images
• 3rd column contains the ground truth segmentation
• 4th column contains Unet result from RGB images (3 channels)
• 5th column contains Unet result from RGB + NIR images (4 channels)
• 6th column contains Unet result from RGB + Green + NIR + NDWI images (6 channels)

Setup

The implementation is developped and successfully run on Ubuntu 20.04; Python 3.9; GPU: RTX 2060 6G VRAM

Please download the weights of MSNet, RTFNet and UNet models with each modality combination from MODEL_CHECKPOINTS and save them in weights folder

To generate the modality combination images such as coir (RGB + NIR), condwi (RGB + NDWI) or cognirndwi (RGB + Green + NIR + NDWI), please following the instruction in prepare_modalities.ipynb notebook in the notebooks folder

To install the dependencies, please use the following

pip3 install -r requirements.txt

Please install Pytorch GPU appropriate with your hardware configuration.

Evaluate

To perform the evaluation of MSNet, RTFNet and UNet models on the WaterMAI's test set by using 3 different modality combinations coir (RGB + NIR), condwi (RGB + NDWI) or cognirndwi (RGB + Green + NIR + NDWI), please following the instruction in evaluate_coir_condwi_cognirndwi.ipynb notebook in the notebooks folder

Prerequisite

1. Please change YOUR_PATH_TO in the notebook to the location in your computer
2. Please run prepare_modalities.ipynb in the notebooks folder for preparing WaterMAI's test set with multiple modality combinations

Training

To perform the training of MSNet, RTFNet and UNet models on the WaterMAI's train set by using 3 different modality combinations coir (RGB + NIR), condwi (RGB + NDWI) or cognirndwi (RGB + Green + NIR + NDWI), please following the instruction in train_coir_condwi_cognirndwi.ipynb notebook in notebooks folder

Prerequisite

1. Please change YOUR_PATH_TO in the notebook to the location in your computer
2. Please run prepare_modalities.ipynb in the notebooks folder for preparing WaterMAI's training set with multiple modality combinations

Reference

Pytorch-UNet

RTFNet

MSNet

Citation

If you find this repo useful for your research, please consider citing our paper:

Ngo, P.L., Pham, V.H., Bui, N.L. et al. Detection of small water bodies for vector control using deep learning on multispectral imagery from unmanned aerial vehicles. Discov Artif Intell 5, 170 (2025). https://doi.org/10.1007/s44163-025-00422-6