Cloud segmentation on satellite images from Sentinel-2 Cloud Mask Catalogue, written in PyTorch. DataLoader is implemented for binary segmentation with CLOUD and CLEAR/CLOUD_SHADOW classes. By default this repo assumes 4 (RGB + NIR) channels, but this can be changed through arguments/config.
pip install -r requirements.txt
or build docker image and run container with:
make build
make
You might have to change base image according to your CUDA and cuDNN versions. If you want to do training, then also change LOCAL_DATASET in Makefile to your dataset path.
Inference is accelerated with ONNX Runtime. You can download model here and then run inference:
python inference.py --model DeepLabV3Plus_resnet101_1678896432.onnx --source data/test_subscene.npy --save --out_folder data/preds/
Predicted masks will be saved in out_folder. You can also plot masks with --show.
Short example of how to do inference in Google Colab is in notebooks/sentinel_segmentation_inference.ipynb
Download and unzip subscenes.zip and masks.zip to your dataset path. This path can be changed in cfg/config.json under dataset_path.
Images can be loaded with different tiling strategies by changing these settings:
- subscene_width, subscene_height - manually resize subscenes and masks before tiling
- train_tile_stride_x, train_tile_stride_y, val_tile_stride_x, val_tile_stride_y - lower stride than tile height/width will lead to overlap
- train_scale, val_scale - automatically scale subscenes and masks so that tiles fit the entire image. This is done after manual resizing with subscene_width/subscene_height. Allowed values are [None, 'down', 'up']
Modify cfg/config.json to your liking and run:
python train.py
You can also override any setting through command-line arguments without modifying the config file, for example:
python3 train.py --epochs 100 --batch_size 128 --lr 0.01 --optimizer AdamW --scheduler cos --warmup_epochs 5 --decoder_name UnetPlusPlus --encoder_name resnet50 --loss CE --no_wandb_log
If you want to utilize more GPUs, enable Distributed Data Parallel (DDP), where --nproc_per_node sets the number of GPUs:
torchrun --standalone --nproc_per_node 2 train.py
Example of training progress:
| Model | mIoU |
|---|---|
| DeepLabV3+ with ResNet101 | 86.43 |
| Unet++ with ResNet101 | 84.55 |
The DeepLabV3+ model with ResNet101 encoder was trained with these settings:
- pretrained encoder on ImageNet
- 9:1 train/val split
- 50 epochs
- 20 epochs early stop
- 64 batch size
- AdamW optimizezr
- 0.0005 learning rate
- 0.005 weight decay
- 3 linear warmup epochs
- linear lr scheduler
- Dice loss
- mean IoU val metric
- augmentations: rotation (max 60 degrees, probability 0.5), horizontal flip (probability 0.5), vertical flip (probability 0.5)
- Multi-class segmentation
- Resume training
- Automated hyperparameter tuning
- Add more metrics
After training, the PyTorch model can be exported to ONNX with:
python export.py --model_file weights/my_model.pt