PyTorch implementation of DRAEM - ICCV2021:

@InProceedings{Zavrtanik_2021_ICCV,
    author    = {Zavrtanik, Vitjan and Kristan, Matej and Skocaj, Danijel},
    title     = {DRAEM - A Discriminatively Trained Reconstruction Embedding for Surface Anomaly Detection},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2021},
    pages     = {8330-8339}
}

To train on the MVtec Anomaly Detection dataset download the data and extract it. The Describable Textures dataset was used as the anomaly source image set in most of the experiments in the paper. You can run the download_dataset.sh script from the project directory to download the MVTec and the DTD datasets to the datasets folder in the project directory:

./scripts/download_dataset.sh

Its requirements.txt file is a mixed file of Conda and Pip, and it's a bit troublesome to install. For the environment, having torch, opencv-python==4.5.3.56, imgaug==0.4.0, and sklearn==0.0 is almost sufficient. If there are still some missing components later, you just need to install one or two more libraries.

Pass the folder containing the training dataset to the train_DRAEM.py script as the --data_path argument and the folder locating the anomaly source images as the --anomaly_source_path argument. The training script also requires the batch size (--bs), learning rate (--lr), epochs (--epochs), path to store checkpoints (--checkpoint_path) and path to store logs (--log_path). Example:

python train_DRAEM.py --gpu_id 0 --obj_id -1 --lr 0.0001 --bs 8 --epochs 700 --data_path ./datasets/mvtec/ --anomaly_source_path ./datasets/dtd/images/ --checkpoint_path ./checkpoints/ --log_path ./logs/

Pretrained DRAEM models for each class of the MVTec anomaly detection dataset are available here. To download the pretrained models directly see ./scripts/download_pretrained.sh.

The pretrained models achieve a 98.1 image-level ROC AUC, 97.5 pixel-wise ROC AUC and a 68.9 pixel-wise AP.

The test script requires the --gpu_id arguments, the name of the checkpoint files (--base_model_name) for trained models, the location of the MVTec anomaly detection dataset (--data_path) and the folder where the checkpoint files are located (--checkpoint_path) with pretrained models can be run with:

python test_DRAEM.py --gpu_id 0 --base_model_name "DRAEM_seg_large_ae_large_0.0001_800_bs8" --data_path ./datasets/mvtec/ --checkpoint_path ./checkpoints/DRAEM_checkpoints/