SAMPolyBuild

This repository is the code implementation of the paper "SAMPolyBuild: Adapting the Segment Anything Model (SAM) for Polygonal Building Extraction" accepted by ISPRS Journal of Photogrammetry and Remote Sensing. Now only the prompt mode is provided, the auto mode will be released soon.

Installation

Conda virtual environment is recommended for installation. Please choose the appropriate version of torch and torchvision according to your CUDA version.

conda create -n sampoly python=3.10 -y
source activate sampoly # or conda activate sampoly
pip install torch==2.0.0+cu117 torchvision==0.15.1+cu117 --index-url https://download.pytorch.org/whl/cu117
pip install -r requirements.txt

Download the SAM vit_b model from here and place it in the 'segment_anything' folder.

Dataset Preparation

SpaceNet Vegas Dataset

We converted the original images of the SpaceNet dataset to 8-bit and the annotations to coco format, and divided them into training, validation, and test sets in the ratio of 8:1:1, which are available for download from here. Place the train, val, test folders in the 'dataset/spacenet' folder.

Custom Dataset

The custom dataset should be in the following format, or change the train_dataset_pth, val_dataset_pth in the train.py and dataset_pth in the test.py to the corresponding path.

dataset
├── dataset_name
    ├── train
    |    ├── images
    |    ├── ann.json
    ├── val
    |    ├── images
    |    ├── ann.json
    ├── test
        ├── images
        ├── ann.json

Training

Single gpu:

python train.py --config configs/prompt_instance_spacenet.json --gpus 0

Multi gpus:

python train.py --config configs/prompt_instance_spacenet.json --gpus 0 1 --distributed

Testing

Evaluate the model on the test set, and save the results:

python test.py

You need to change the --task_name to the corresponding training task name, and the other arguments will be set automatically according to training configuration.

If you want to use our trained model to evaluate, you can download prompt_instance_spacenet.ckpt (extract code: be8m) and change the following code in the test.py:

args = load_args(parser,path='configs/prompt_instance_spacenet.json')
args.checkpoint = 'prompt_instance_spacenet.ckpt'

Acknowledgement

This project is developed based on the Segment Anything Model (SAM) project. We thank the authors for their contributions.

Citation

If you use the code of this project in your research, please refer to the bibtex below to cite SAMPolyBuild.

@article{wang2024sampolybuild,
  title={SAMPolyBuild: Adapting the Segment Anything Model for polygonal building extraction},
  author={Wang, Chenhao and Chen, Jingbo and Meng, Yu and Deng, Yupeng and Li, Kai and Kong, Yunlong},
  journal={ISPRS Journal of Photogrammetry and Remote Sensing},
  volume={218},
  pages={707--720},
  year={2024},
  publisher={Elsevier},
  doi = {10.1016/j.isprsjprs.2024.09.018}
}

License

This project is licensed under the Apache 2.0 license.

Contact

If you have any questions, please contact [email protected].