Official codebase for "Segment Anything Model for Road Network Graph Extraction", CVPRW 2024

https://arxiv.org/pdf/2403.16051.pdf

The paper has been accepted by IEEE/CVF Computer Vision and Pattern Recognition Conference (CVPR) 2024, 2nd Workshop on Scene Graphs and Graph Representation Learning.

Received the best paper award of the workshop. Big thanks to the organizers for the recognition!!

Demos

Predicted road network graph in a large region (2km x 2km).

Predicted road network graphs and corresponding masks in dense urban with complex and irregular structures.

Installation

You need the following:

• an Nvidia GPU with latest CUDA and driver.
• the latest pytorch.
• pytorch lightning.
• wandb.
• Go, just for the APLS metric (we should really re-write this with pure python when time allows).
• and pip install whatever is missing.

Getting Started

SAM Preparation

Download the ViT-B checkpoint from the official SAM directory. Put it under:
-sam_road
--sam_ckpts
---sam_vit_b_01ec64.pth

Data Preparation

Refer to the instructions in the RNGDet++ repo (https://github.com/TonyXuQAQ/RNGDetPlusPlus) to download City-scale and SpaceNet datasets. Put them in the main directory, structure like:
-sam_road
--cityscale
---20cities
--spacenet
---RGB_1.0_meter

Download links copied from https://github.com/TonyXuQAQ/RNGDetPlusPlus

SpaceNet

https://drive.google.com/uc?id=1FiZVkEEEVir_iUJpEH5NQunrtlG0Ff1W

The data_split.json is copied from the dataset.json in this folder.

CityScale

https://drive.google.com/uc?id=1R8sI1RmFe3rUfWMQaOfsYlBDHpQxFH-H

Find the 20cities folder under this folder.

Then, run "python generate_labes.py" under both dirs.

Training

City-scale dataset:
python train.py --config=config/toponet_vitb_512_cityscale.yaml

SpaceNet dataset:
python train.py --config=config/toponet_vitb_256_spacenet.yaml

You can find the checkpoints under lightning_logs dir.

Inference

python inferencer.py --config=path_to_the_same_config_for_training --checkpoint=path_to_ckpt
This saves the inference results and visualizations.

Inferencing with our checkpoints:

Cityscale:

python inferencer.py --config=config/toponet_vitb_512_cityscale.yaml --checkpoint=/path_to/cityscale_vitb_512_e10.ckpt

Spacenet:

python inferencer.py --config=config/toponet_vitb_256_spacenet.yaml --checkpoint=/path_to/spacenet_vitb_256_e10.ckpt

Test

Go to cityscale_metrics or spacenet_metrics, and run
bash eval_schedule.bash

Check that script for details. It runs both APLS and TOPO and stores scores to your output dir.

Our Checkpoints

congrui/sam_road

Citation

@article{hetang2024segment,
  title={Segment Anything Model for Road Network Graph Extraction},
  author={Hetang, Congrui and Xue, Haoru and Le, Cindy and Yue, Tianwei and Wang, Wenping and He, Yihui},
  journal={arXiv preprint arXiv:2403.16051},
  year={2024}
}

Acknowledgement

We sincerely appreciate the authors of the following codebases which made this project possible:

• Segment Anything Model
• RNGDet++
• SAMed
• Detectron2

TODO List

• Basic instructions
• Organize configs
• Add dependency list
• Add demos
• Add trained checkpoints