Repository files navigation

This is an official implementation of “Registration based Few-Shot Anomaly Detection” (RegAD) with PyTorch, accepted by ECCV 2022 (Oral).

Paper Link

@inproceedings{huang2022regad,  title={Registration based Few-Shot Anomaly Detection}  author={Huang, Chaoqin and Guan, Haoyan and Jiang, Aofan and Zhang, Ya and Spratlin, Michael and Wang, Yanfeng},  booktitle={European Conference on Computer Vision (ECCV)},  year={2022}}

Abstract: This paper considers few-shot anomaly detection (FSAD), a practical yet under-studied setting for anomaly detection (AD), where only a limited number of normal images are provided for each category at training. So far, existing FSAD studies follow the one-model-per-category learning paradigm used for standard AD, and the inter-category commonality has not been explored. Inspired by how humans detect anomalies, i.e., comparing an image in question to normal images, we here leverage registration, an image alignment task that is inherently generalizable across categories, as the proxy task, to train a category-agnostic anomaly detection model. During testing, the anomalies are identified by comparing the registered features of the test image and its corresponding support (normal) images. As far as we know, this is the first FSAD method that trains a single generalizable model and requires no re-training or parameter fine-tuning for new categories.

Keywords: Anomaly Detection, Few-Shot Learning, Registration

• python >= 3.7.11
• pytorch >= 1.11.0
• torchvision >= 0.12.0
• numpy >= 1.19.5
• scipy >= 1.7.3
• skimage >= 0.19.2
• matplotlib >= 3.5.2
• kornia >= 0.6.5
• tqdm

1. Download the MVTec datasethere.
2. Download the support dataset for few-shot anomaly detection onGoogle Drive orBaidu Disk (i9rx)and unzip the dataset. For those who have problem downloading the support set, please optional download categories of capsule and grid onBaidu Disk (pll9) andBaidu Disk (ns0n).

   tar -xvf support_set.tar

   We hope the followers could use these support datasets to make a fair comparison between different methods.
3. Download the pre-train models onGoogle Drive orBaidu Disk (4qyo)and unzip the checkpoint files.

   tar -xvf save_checkpoints.tar

After the preparation work, the whole project should have the following structure:

./RegAD├── README.md├── train.py                                  # training code├── test.py                                   # testing code├── MVTec                                     # MVTec dataset files│   ├── bottle│   ├── cable│   ├── ...                  │   └── zippper├── support_set                               # MVTec support dataset files│   ├── 2│   ├── 4                 │   └── 8├── models                                    # models and backbones│   ├── stn.py  │   └── siamese.py├── losses                                    # losses│   └── norm_loss.py  ├── datasets                                  # dataset                      │   └── mvtec.py├── save_checkpoints                          # model checkpoint files                  └── utils                                     # utils    ├── utils.py    └── funcs.py

pythontest.py--obj $target-object--shot $few-shot-number--stn_moderotation_scale

For example, if run on the categorybottle withk=2:

pythontest.py--objbottle--shot2--stn_moderotation_scale

pythontrain.py--obj $target-object--shot $few-shot-number--data_typemvtec--data_path ./MVTec/--epochs50--batch_size32--lr0.0001--momentum0.9--inferences10--stn_moderotation_scale

For example, to train a RegAD model on the MVTec dataset onbottle withk=2, simply run:

pythontrain.py--objbottle--shot2--data_typemvtec--data_path ./MVTec/--epochs50--batch_size32--lr0.0001--momentum0.9--inferences10--stn_moderotation_scale

Then you can run the evaluation using:

pythontest.py--objbottle--shot2--stn_moderotation_scale

Results of few-shot anomaly detection and localization with k=2:

Results of few-shot anomaly detection and localization with k=4:

Results of few-shot anomaly detection and localization with k=8:

We borrow some codes fromSimSiam,STN andPaDiM

If you have any problem with this code, please feel free to contacthuangchaoqin@sjtu.edu.cn.