This is the official code repository for "Wavelet and Prototype Augmented Query-based Transformer for Pixel-level Surface Defect Detection." The paper could be found at Link

:fire: News

2025-12-25: We release the results of WPFormer trained on COD and SOD datasets.

1. Requirements

• python==3.7.13
• cudatoolkit==11.3.1
• pytorch==1.11.0

2. Results on COD

3. Results on SOD

4. Get Start

1. Download Datasets and Checkpoints.

• Datasets:

By default, you can put datasets into the folder 'Dataset'.

• Checkpoints:

By default, you can put pretrained backbone checkpoints into the folder 'model' and modify the model path in "WPFormer.py".

2. Test.

modify the test dataset and model path in "defect_test.py".

run defect_test.py

3. Eval.

The metric evaluation code "sod_metrics.py" is sourced from PySODMetrics. By default, you can download prediction maps and unzip it into the main folder, and modify the dataset path and prediction maps path in "eval.py".

run eval.py

5. Pretrained models and results

All predictions obtained by different methods on five datasets are provided below:

Dataset Name	Dataset Download	Backbone	Input size	Config	Checkpoints	Prediction maps
ESDIs-SOD	Link	PVTV2-B2	384x384	channel=64, bs=8, lr=8e-5, epoch=150	Link	Link
CrackSeg9k	Link	PVTV2-B2	384x384	channel=64, bs=4, lr=8e-5, epoch=60	Link	Link
ZJU-Leaper	Link	PVTV2-B2	384x384	channel=64, bs=4, lr=8e-5, epoch=24	Link	Link
SOD	Link	PVTV2-B4	384x384	channel=128, bs=6, lr=4e-5, epoch=60	Link	Link
COD	Link	PVTV2-B4	384x384	channel=128, bs=16, lr=4e-5, epoch=150	Link	Link
			512x512	channel=128, bs=8, lr=4e-5, epoch=150	Link	Link

Citation

@inproceedings{yan2025wavelet,
  title={Wavelet and Prototype Augmented Query-based Transformer for Pixel-level Surface Defect Detection},
  author={Yan, Feng and Jiang, Xiaoheng and Lu, Yang and Cao, Jiale and Chen, Dong and Xu, Mingliang},
  booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
  pages={23860--23869},
  year={2025}
}

Acknowledgement

We would like to acknowledge the contributions of public projects, such as MaskFormer, Mask2Former, whose code has been utilized in this repository.