A Pytorch implementation of several semantic segmentation methods on the dataset introduced in the paper Learning Aerial Image Segmentation from Online Maps.
Create a virtual environment and install the dependencies:
python3 -m venv env
source env/bin/activate
pip3 install -r requirements.txt
Download the dataset:
cd dataset && sh download_dataset.sh
Examples from the dataset:
UNet, FastSCNN and DeeplabV3 are implemented.
python3 aerial-segmentation/train.py
optional arguments:
-h, --help show this help message and exit
--model {UNet,Deeplabv3,FastSCNN}
Network model to be trained (default: UNet)
--loss {FocalLoss,DiceLoss,CrossEntropyLoss}
Loss function (default: FocalLoss)
--optimizer {SGD,Adam}
Optimizer (default: Adam)
--resample-size {0,1,2,3,4,5}
Number of crops to be used for each image. If 5 is selected, all the 4 corner crops and 1 center crop will be added as augmentation (default: 5)
--batch-coeff BATCH_COEFF
Batch size is equal to [batch_coeff] x [resample_size] (default: 1)
--lr LR Learning rate (default: 1e-3)
--epochs EPOCHS Maximum number of epochs (default: 50)
--image-size IMAGE_SIZE
Image size (default: 256)
- Implement argparse for different training options
- Add evaluation results
- Add Docker support
- Deploy in C++