This is the code to the WACV 2023 paper "Asymmetric Student-Teacher Networks for Industrial Anomaly Detection" by Marco Rudolph, Tom Wehrbein, Bodo Rosenhahn and Bastian Wandt.
You will need Python 3.7.7 and the packages specified in requirements.txt. We recommend setting up a virtual environment with pip and installing the packages there.
Install packages with:
$ pip install -r requirements.txt
All configurations concerning data, model, training, visualization etc. can be made in config.py. The default configuration will run a training with paper-given parameters for MVTec 3D-AD.
The following steps guide you from your dataset to your evaluation:
- Set your dataset_dir in config.py. This is the directory which contains the subdirectories of the classes you want to process. It should be configured whether the datasets contains 3D scans (set use_3D_dataset in config.py). If 3D data is present (use_3D_dataset=True), the data structure from MVTec 3D-AD is assumed. It can be chosen which data domain should be used for detection (set mode in config.py). If the dataset does not contain 3D data (use_3D_dataset=False), the data structure from the MVTec AD dataset is assumed. Both dataset structures are described in the documentation of load_img_datasets in utils.py.
- preprocessing.py: It is recommended to pre-extract the features beforehand to save training time. This script also preprocesses 3D scans. Alternatively, the raw images are used in training when setting pre_extracted=False in config.py.
- train_teacher.py: Trains the teacher and saves the model to models/...
- train_student.py: Trains the student and saves the model to models/...
- eval.py: Evaluates the student-teacher-network (image-level results and localization/segmentation). Additionally, it creates ROC curves, anomaly score histograms and localization images.
Note: Due to a bug in the original implementation for the RGB+3D setting the performance on MVTec 3D-AD is about 2% better than originally reported.
Some code of an old version of the FrEIA framework was used for the implementation of Normalizing Flows. Follow their tutorial if you need more documentation about it.
Please cite our paper in your publications if it helps your research. Even if it does not, you are welcome to cite us.
@inproceedings { RudWeh2023,
author = {Marco Rudolph and Tom Wehrbein and Bodo Rosenhahn and Bastian Wandt},
title = {Asymmetric Student-Teacher Networks for Industrial Anomaly Detection},
booktitle = {Winter Conference on Applications of Computer Vision (WACV)},
year = {2023},
month = jan
}
This project is licensed under the MIT License.