SGAT4PASS

Introduction

This is the official implementation of the paper SGAT4PASS: Spherical Geometry-Aware Transformer for PAnoramic Semantic Segmentation (IJCAI 2023)[PDF].

Environments

conda create -n SGAT4PASS python=3.7.7
conda activate SGAT4PASS
cd ~/path/to/SGAT4PASS 
conda install pytorch==1.8.1 torchvision==0.9.1 torchaudio==0.8.1 cudatoolkit=11.1 -c pytorch -c conda-forge
pip install -r requirements.txt
python setup.py develop --user
# Optional: install apex follow: https://github.com/NVIDIA/apex

Data Preparation

Prepare datasets:

Stanford2D3D

datasets/
├── Stanford2D3D
│   ├── area_1
│   ├── area_2
│   ├── area_3
│   ├── area_4
│   ├── area_5a
│   ├── area_5b
│   └── area_6

Prepare pretrained weights, which can be found in the public repository of SegFormer.

pretrained/
├── mit_b1.pth
└── mit_b2.pth

Train

For example, to use 4 A100 GPUs to run the experiments on Stanford2D3D dataset:

LR=0.00008
EPOCHS=150
PER_PIXEL_WEIGHT=0.3 # Panorama-Aware Loss Weight
OFFSET_WEIGHT=0.3 # SDPE Loss Weight
REPROJECTION=True # whether to use reprojection
X_MAX=10 # the maximum value of x-axis
Y_MAX=10 # the maximum value of y-axis
Z_MAX=360 # the maximum value of z-axis
FOLD=1 # the fold of the Stanford2D3D dataset

python -m torch.distributed.launch --nproc_per_node=4 --master_port 30005 tools/train_s2d3d_span.py \
    --config-file configs/stanford2d3d_pan/SGAT4PASS_small_1080x1080_line_axis_a100_xyz_mask_loss.yaml \
    TRAIN.MODEL_SAVE_DIR workdirs/SGAT4PASS_small_1080x1080_line_${OFFSET_WEIGHT}_axis_${OFFSET_WEIGHT}_mask_loss_${PER_PIXEL_WEIGHT}_lr${LR}_epoch_${EPOCHS} \
    SOLVER.LR $LR \
    TRAIN.EPOCHS $EPOCHS \
    SOLVER.PER_PIXEL_WEIGHT $PER_PIXEL_WEIGHT \
    SOLVER.OFFSET_LINE_WEIGHT $OFFSET_WEIGHT \
    SOLVER.OFFSET_AXIS_WEIGHT $OFFSET_WEIGHT \
    AUG.REPROJECTION $REPROJECTION \
    DATASET.FOLD $FOLD \
    AUG.X_MAX $X_MAX \
    AUG.Y_MAX $Y_MAX \
    AUG.Y_MAX $Z_MAX

Test

Download the models from GoogleDrive and put them in ./checkpoints/ folder.

./checkpoints
├── Stanford2D3D_Fold_1
│   ├── best_model_fold_1.pth
├── Stanford2D3D_Fold_2
│   ├── best_model_fold_2.pth
└── Stanford2D3D_Fold_3
    ├── best_model_fold_3.pth

Traditional Testing

python -m torch.distributed.launch --nproc_per_node=4 --master_port 30005 tools/eval_s2d3d_span.py \
    --config-file configs/stanford2d3d_pan/SGAT4PASS_small_1080x1080_Fold_1.yaml \
    TEST.TEST_MODEL_PATH checkpoints/Stanford2D3D_Fold_3/best_model_fold_1.pth \
    DATASET.FOLD 1

python -m torch.distributed.launch --nproc_per_node=4 --master_port 30005 tools/eval_s2d3d_span.py \
    --config-file configs/stanford2d3d_pan/SGAT4PASS_small_1080x1080_Fold_2.yaml \
    TEST.TEST_MODEL_PATH checkpoints/Stanford2D3D_Fold_3/best_model_fold_2.pth \
    DATASET.FOLD 2

python -m torch.distributed.launch --nproc_per_node=4 --master_port 30005 tools/eval_s2d3d_span.py \
    --config-file configs/stanford2d3d_pan/SGAT4PASS_small_1080x1080_Fold_3.yaml \
    TEST.TEST_MODEL_PATH checkpoints/Stanford2D3D_Fold_3/best_model_fold_3.pth \
    DATASET.FOLD 3

SGA Testing

First, change the value of the cfg.TEST.ROTATIONS to the angle to be tested in the segmentron/config/settings.py.

Then, run bash as follows

python -m torch.distributed.launch --nproc_per_node=4 --master_port 30005 tools/eval_s2d3d_span.py \
    --config-file configs/stanford2d3d_pan/SGAT4PASS_small_1080x1080_Fold_1.yaml \
    TEST.TEST_MODEL_PATH checkpoints/Stanford2D3D_Fold_1/best_model_fold_1.pth \
    TEST.SGA True \
    DATASET.FOLD 1

Note that TEST.SGA must be set to True

References

Our repository is heavily based on Trans4PASS

Thank them for their excellent work!

License

This repository is under the Apache-2.0 license. For commercial use, please contact with the authors.

Citations

If you are interested in this work, please cite the following works:

SGAT4PASS [PDF]

@article{li2023sgat4pass,
  title={SGAT4PASS: Spherical Geometry-Aware Transformer for PAnoramic Semantic Segmentation},
  author={Li, Xuewei and Wu, Tao and Qi, Zhongang and Wang, Gaoang and Shan, Ying and Li, Xi},
  journal={arXiv preprint arXiv:2306.03403},
  year={2023}
}

TencentARC/SGAT4PASS