Skip to content

Official implementation of ScoreSeg: Leveraging Score-based Generative model for Self-Supervised Semantic Segmentation of Remote Sensing

License

Notifications You must be signed in to change notification settings

moonbow721/ScoreSeg

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

6 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

ScoreSeg

🎉 Our paper "ScoreSeg: Leveraging Score-based Generative Model for Self-Supervised Semantic Segmentation of Remote Sensing" has been accepted by IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing. Check out the paper on IEEE Xplore.

Official implementation of ScoreSeg: Leveraging Score-based Generative model for Self-Supervised Semantic Segmentation of Remote Sensing.

Environment

  • This repository is built and tested in RTX 3090Ti GPUs, CUDA==11.4.

  • We recommend you to use Anaconda to create a conda environment:

    conda create -n ScoreSeg python=3.9
    conda activate ScoreSeg
  • Then, install the corresponding Pytorch version

    conda install pytorch=1.12.1 torchvision=0.13.1 cudatoolkit=11.3 -c pytorch
  • Other requirements

    pip install -r requirements.txt
  • CUDA operators for Deformable Transformer

    cd ./model/deformalbe_ops
    sh ./make.sh

Data Preprocessing

We follow the folder structure and data preprocessing in the GeoSeg. As a result, for example, 13824 tiles of 256 x 256 crops can be obtained in the Potsdam dataset for training while 8064 tiles for testing.

Note that offline data augmentation should be avoided since the following sampling procedure may collect more images.

You can randomly select portions of the dataset, e.g., 1%, 5% and 10% in our experiments. For example,

python dataset-split.py --dataset_name potsdam --phase train

Training score-based generative models

We recommend you to download the pretrained DDPM released in ddpm-cd, which can save your time to pretrain a score-based generative model (or so-called diffusion model) on remote sensing images yourself.

Or, alternatively, you can also run the following command to start the training procedure on 4 GPUs in the DistributedDataParallel (DDP) style.

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun  --nnodes 1 --master_port 29500 \
--nproc_per_node=4 \ 
ScoreTrain.py --config config/score_pretraining.json --phase train

It will take around 2 days to train a score-based generative model on mixture of Potsdam and Deepglobe datasets.

Training and testing segmentation modules

After the pretraining of score-based generative models, you can specify the load path of its pth file in the segmentation config files, e.g., potsdam_segmentation.json

Then run the following command to train the segmentation modules on 2 GPUs and enable wandb (optional) for logging.

CUDA_VISIBLE_DEVICES=0,1 torchrun  --nnodes 1 --master_port 29500 \
--nproc_per_node=2 \
ScoreSeg.py --config config/potsdam_segmentation.json -enable_wandb

This will record checkpoints and visualization of segmentation maps in the specified paths in config files automatically.

For testing, remember to specify the 'resume_state' parameter in config files first.

CUDA_VISIBLE_DEVICES=0,1 torchrun  --nnodes 1 --master_port 29500 \
--nproc_per_node=2 \
ScoreSeg.py --config config/potsdam_segmentation.json -enable_wandb  --phase test

Pretrained weights

For convenience, we provide several pretrained segmentation modules weights here.

Dataset Ratio OA mF1 mIoU Url
Potsdam 1% 74.62 69.84 55.60 link
Vaihingen 10% 83.75 67.39 56.33 link
Deepglobe 5% 78.66 66.68 53.68 link

Note:

  1. These models are all trained with the score-based models weights in ddpm-cd.
  2. These models use the same hyparameters as provided config examples.
  3. "Ratio" means the proportion of dataset labels used.

References

The code of score-based models is from ddpm-cd. The implementation of deformable transformer and segmentation losses are from Deformable-DETR and GeoSeg respectively.

Citation

If you find this work useful for your research, please consider citing:

@article{lu2023scoreseg,
  title={ScoreSeg: Leveraging Score-based Generative Model for Self-Supervised Semantic Segmentation of Remote Sensing},
  author={Lu, Junzhe and He, Guangjun and Dou, Hongkun and Gao, Qing and Fang, Leyuan and Deng, Yue},
  journal={IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing},
  year={2023},
  publisher={IEEE}
}

About

Official implementation of ScoreSeg: Leveraging Score-based Generative model for Self-Supervised Semantic Segmentation of Remote Sensing

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 79.3%
  • Cuda 18.6%
  • C++ 1.9%
  • Shell 0.2%