Official PyTorch implementation of DeSTSeg - CVPR 2023

We use the MVTec AD dataset for experiments. To simulate anomalous image, the Describable Textures Dataset (DTD) is also adopted in our work. Users can run the download_dataset.sh script to download them directly.

./scripts/download_dataset.sh

Please install the dependency packages using the following command by pip:

pip install -r requirements.txt

To get started, users can run the following command to train the model on all categories of MVTec AD dataset:

python train.py --gpu_id 0 --num_workers 16

Users can also customize some default training parameters by resetting arguments like --bs, --lr_DeST, --lr_res, --lr_seghead, --steps, --DeST_steps, --eval_per_steps, --log_per_steps, --gamma and --T.

To specify the training categories and the corresponding data augmentation strategies, please add the argument --custom_training_category and then add the categories after the arguments --no_rotation_category, --slight_rotation_category and --rotation_category. For example, to train the screw category and the tile category with no data augmentation strategy, just run the following command:

python train.py --gpu_id 0 --num_workers 16 --custom_training_category --no_rotation_category screw tile

To test the performance of the model, users can run the following command:

python eval.py --gpu_id 0 --num_workers 16

Download pretrained checkpoints here and put the checkpoints under <project_dir>/saved_model/.

@inproceedings{zhang2023destseg,
  title={DeSTSeg: Segmentation Guided Denoising Student-Teacher for Anomaly Detection},
  author={Zhang, Xuan and Li, Shiyu and Li, Xi and Huang, Ping and Shan, Jiulong and Chen, Ting},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={3914--3923},
  year={2023}
}