LMA-Swin

Introduction

In this repository, we offer the code of LMA-Swin, a semantic segmentation model based on PyTorch which combines Convolutional Neural Network (CNN) features with Transformer features through modulation, addressing the limitations of Transformer-based models in preserving fine-grained details and local edge segmentation. You can download the trained model weights on BaiduDisk. Code: lmas

Network

Dependency Installation

conda create -n airs python=3.8
conda activate airs
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install -r GeoSeg/requirements.txt

Data Preparation

Please download the datasets from the official website and prepare the data following the following file structure:

airs
├── GeoSeg (code)
├── pretrain_weights (pretrained weights of backbones)
├── model_weights (save the model weights trained on datasets)
├── fig_results (save the masks predicted by models)
├── lightning_logs (CSV format training logs)
├── data
│   ├── LoveDA
│   │   ├── Train
│   │   │   ├── Urban
│   │   │   │   ├── images_png (original images)
│   │   │   │   ├── masks_png (original masks)
│   │   │   │   ├── masks_png_convert (converted masks used for training)
│   │   │   │   ├── masks_png_convert_rgb (original rgb format masks)
│   │   │   ├── Rural
│   │   │   │   ├── images_png 
│   │   │   │   ├── masks_png 
│   │   │   │   ├── masks_png_convert
│   │   │   │   ├── masks_png_convert_rgb
│   │   ├── Val (the same with Train)
│   │   ├── Test
│   │   ├── train_val (Merge Train and Val)
│   ├── uavid
│   │   ├── uavid_train (original)
│   │   ├── uavid_val (original)
│   │   ├── uavid_test (original)
│   │   ├── uavid_train_val (Merge uavid_train and uavid_val)
│   │   ├── train (processed)
│   │   ├── val (processed)
│   │   ├── train_val (processed)
│   ├── vaihingen
│   │   ├── train_images (original)
│   │   ├── train_masks (original)
│   │   ├── test_images (original)
│   │   ├── test_masks (original)
│   │   ├── test_masks_eroded (original)
│   │   ├── train (processed)
│   │   ├── test (processed)
│   ├── potsdam (the same with vaihingen)

Training

python GeoSeg/train_supervision.py -c GeoSeg/config/vaihingen/lmaswin.py

"-c" means the path of the config, use different config to train different models.

Testing

Vaihingen

python GeoSeg/vaihingen_test.py -c GeoSeg/config/vaihingen/lmaswin.py -o fig_results/vaihingen/lmaswin --rgb -t 'd4'

Potsdam

python GeoSeg/potsdam_test.py -c GeoSeg/config/potsdam/lmaswin.py -o fig_results/potsdam/lmaswin --rgb -t 'lr'

"-c" denotes the path of the config, Use different config to test different models

"-o" denotes the output path

"-t" denotes the test time augmentation (TTA), can be [None, 'lr', 'd4'], default is None, 'lr' is flip TTA, 'd4' is multiscale TTA

"--rgb" denotes whether to output masks in RGB format

Visualization Results

Acknowledgement

Many thanks for the GeoSeg's contributions to LMA-Swin.