[ACM MM2024] Learning Spectral-Decomposed Tokens for Domain Generalized Semantic Segmentation

Jingjun Yi¹, Qi Bi², Hao Zheng³, Haolan Zhan⁴, Wei Ji⁵, Yawen Huang³, Yuexiang Li⁶, Yefeng Zheng^3
1 Wuhan University ² University of Amsterdam ³ Tencent Youtu Lab ⁴ Monash University ⁵ Yale University ⁶ Guangxi Medical University

Project page: https://github.com/JingjunYi/SET

Paper: https://arxiv.org/abs/2407.18568

Spectral-dEcomposed Token (SET) is a novel framework designed to enhance domain generalized semantic segmentation by decomposing frozen VFM features into phase (content) and amplitude (style) components in the frequency space, allowing learnable tokens to process these components separately. This approach bridges the gap between style variations and static tokens, resulting in state-of-the-art performance in DGSS benckmarks.

Visualization

Trained on Cityscapes, SET generalizes to unseen domains: BDD, Mapillary, GTAV, Synthia.

Performance Under Various Settings (DINOv2).

Environment Setup

To set up your environment, execute the following commands:

conda create -n set -y
conda activate set
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia -y
pip install -U openmim
mim install mmengine
mim install "mmcv>=2.0.0"
pip install "mmsegmentation>=1.0.0"
pip install "mmdet>=3.0.0"
pip install xformers=='0.0.20' # optional for DINOv2
pip install -r requirements.txt
pip install future tensorboard

Dataset Preparation

The Preparation is similar as our former work CMFormer. Specifically, all the data is converted to the Cityscapes format.

Pretraining Weights

• Download: Download pre-trained weights from facebookresearch for testing. Place them in the project directory without changing the file name.
• Convert: Convert pre-trained weights for training or evaluation.

  python tools/convert_models/convert_dinov2.py checkpoints/dinov2_vitl14_pretrain.pth checkpoints/dinov2_converted.pth

  (optional for 1024x1024 resolution)

  python tools/convert_models/convert_dinov2.py checkpoints/dinov2_vitl14_pretrain.pth checkpoints/dinov2_converted_1024x1024.pth --height 1024 --width 1024

Evaluation

Run the evaluation:

  python tools/test.py configs/my/citys_rein_dinov2_mask2former_512x512_bs1x4.py exps/exp0429_syn/iter_40000.pth --backbone checkpoints/dinov2_converted.pth

Training

Run the training

python tools/train.py configs/my/citys_rein_dinov2_mask2former_512x512_bs1x4.py --work-dir exps/exp0322

Acknowledgment

Our implementation is primarily based on the following repositories, with significant influence from Rein. Thanks for their authors.

• MMSegmentation
• Rein

Contact

For further information or questions, please contact Jingjun Yi via rsjingjuny@whu.edu.cn or Qi Bi via q.bi@uva.nl.

Citation

If you find our code or data helpful, please cite our paper:

@article{yi2024learning,
  title={Learning Spectral-Decomposed Tokens for Domain Generalized Semantic Segmentation},
  author={Yi, Jingjun and Bi, Qi and Zheng, Hao and Zhan, Haolan and Ji, Wei and Huang, Yawen and Li, Yuexiang and Zheng, Yefeng},
  journal={arXiv preprint arXiv:2407.18568},
  year={2024}
}