Awesome

✨MuSc (ICLR 2024)✨

This is an official PyTorch implementation for "MuSc : Zero-Shot Industrial Anomaly Classification and Segmentation with Mutual Scoring of the Unlabeled Images" (MuSc)

Authors: Xurui Li^1* | Ziming Huang^1* | Feng Xue³ | Yu Zhou^1,2

Institutions: ¹Huazhong University of Science and Technology | ²Wuhan JingCe Electronic Group Co.,LTD | ³University of Trento

🧐 Arxiv | OpenReview

📖 Chinese README

<a href='#all_catelogue'>Go to Catalogue</a>

🙈TODO list:

• ⬜️ Using some strategies to reduce the inference time per image from 955.3ms to 249.8ms.
• ⬜️ Compatibility with more industrial datasets.
• ⬜️ Compatibility with more visual backbones, e.g. Vision Mamba.

📣Updates:

04/11/2024

1. The comparisons with the zero/few-shot methods in CVPR 2024 have been added to <a href='#compare_sota'>Compare with SOTA k-shot Methods.</a>
2. Fixed some bugs in models/backbone/_backbones.py.

03/22/2024

1. The supported codes for BTAD dataset are provided.
2. Some codes are modified to support larger batch_size.
3. Some codes are optimized to obtain faster speeds.
4. <a href='#results_backbones'>Results of different backbones</a> in MVTec AD, VisA and BTAD datasets are provided.
5. <a href='#results_datasets'>The detailed results of different datasets</a> are provided.
6. <a href='#inference_time'>The inference time of different backbones</a> is provided.
7. <a href='#compare_sota'>The comparisons with SOTA zero/few-shot methods</a> are provided. This table will be updated continuously.
8. We summarize the <a href='#FAQ'> frequently asked questions </a> from users when using MuSc, and give the answers.
9. We add README in Chinese.

02/01/2024

Initial commits:

1. The complete code of our method MuSc in paper is released.
2. This code is compatible with image encoder (ViT) of CLIP and ViT pre-trained with DINO/DINO_v2.

🎖️Compare with SOTA k-shot methods <a href='#all_catelogue'>[Go to Catalogue]</a>

We will continuously update the following table to compare our MuSc with the newest zero-shot and few-shot methods. "-" indicates that the authors did not measure this metric in their paper.

MVTec AD

VisA

📖Catalogue

• <a href='#abstract'>1. Abstract</a>
• <a href='#setup'>2. Environment setup</a>
• <a href='#datasets'>3. Datasets download</a>
  • <a href='#datatets_mvtec_ad'>MVTec AD</a>
  • <a href='#datatets_visa'>VisA</a>
  • <a href='#datatets_btad'>BTAD</a>
• <a href='#run_musc'>4. Run MuSc</a>
• <a href='#rscin'>5. Run RsCIN</a>
• <a href='#results_datasets'>6. Results of different datasets</a>
• <a href='#results_backbones'>7. Results of different backbones</a>
• <a href='#inference_time'>8. Inference time</a>
• <a href='#FAQ'>9. Frequently Asked Questions</a>
• <a href='#citation'>10. Citation</a>
• <a href='#thanks'>11. Thanks</a>
• <a href='#license'>12. License</a>

👇Abstract: <a href='#all_catelogue'>[Back to Catalogue]</a>

This paper studies zero-shot anomaly classification (AC) and segmentation (AS) in industrial vision. We reveal that the abundant normal and abnormal cues implicit in unlabeled test images can be exploited for anomaly determination, which is ignored by prior methods. Our key observation is that for the industrial product images, the normal image patches could find a relatively large number of similar patches in other unlabeled images, while the abnormal ones only have a few similar patches.

We leverage such a discriminative characteristic to design a novel zero-shot AC/AS method by Mutual Scoring (MuSc) of the unlabeled images, which does not need any training or prompts. Specifically, we perform Local Neighborhood Aggregation with Multiple Degrees (LNAMD) to obtain the patch features that are capable of representing anomalies in varying sizes. Then we propose the Mutual Scoring Mechanism (MSM) to leverage the unlabeled test images to assign the anomaly score to each other. Furthermore, we present an optimization approach named Re-scoring with Constrained Image-level Neighborhood (RsCIN) for image-level anomaly classification to suppress the false positives caused by noises in normal images.

The superior performance on the challenging MVTec AD and VisA datasets demonstrates the effectiveness of our approach. Compared with the state-of-the-art zero-shot approaches, MuSc achieves a $\textbf{21.1}$% PRO absolute gain (from 72.7% to 93.8%) on MVTec AD, a $\textbf{19.4}$% pixel-AP gain and a $\textbf{14.7}$% pixel-AUROC gain on VisA. In addition, our zero-shot approach outperforms most of the few-shot approaches and is comparable to some one-class methods.

😊Compare with other 0-shot methods

😊Compare with other 4-shot methods

🎯Setup: <a href='#all_catelogue'>[Back to Catalogue]</a>

Environment:

• Python 3.8
• CUDA 11.7
• PyTorch 2.0.1

Clone the repository locally:

git clone https://github.com/xrli-U/MuSc.git

Create virtual environment:

conda create --name musc python=3.8
conda activate musc

Install the required packages:

pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
pip install -r requirements.txt

👇Datasets Download: <a href='#all_catelogue'>[Back to Catalogue]</a>

Put all the datasets in ./data folder.

MVTec AD

data
|---mvtec_anomaly_detection
|-----|-- bottle
|-----|-----|----- ground_truth
|-----|-----|----- test
|-----|-----|----- train
|-----|-- cable
|-----|--- ...

VisA

data
|----visa
|-----|-- split_csv
|-----|-----|--- 1cls.csv
|-----|-----|--- ...
|-----|-- candle
|-----|-----|--- Data
|-----|-----|-----|----- Images
|-----|-----|-----|--------|------ Anomaly 
|-----|-----|-----|--------|------ Normal 
|-----|-----|-----|----- Masks
|-----|-----|-----|--------|------ Anomaly 
|-----|-----|--- image_anno.csv
|-----|-- capsules
|-----|--- ...

VisA dataset need to be preprocessed to separate the train set from the test set.

python ./datasets/visa_preprocess.py

BTAD

data
|---btad
|-----|--- 01
|-----|-----|----- ground_truth
|-----|-----|----- test
|-----|-----|----- train
|-----|--- 02
|-----|--- ...

💎Run MuSc: <a href='#all_catelogue'>[Back to Catalogue]</a>

We provide two ways to run our code.

python

python examples/musc_main.py

Follow the configuration in ./configs/musc.yaml.

script

sh scripts/musc.sh

The configuration in the script musc.sh takes precedence.

The key arguments of the script are as follows:

• --device: GPU_id.
• --data_path: The directory of datasets.
• --dataset_name: Dataset name.
• --class_name: Category to be tested. If the parameter is set to ALL, all the categories are tested.
• --backbone_name: Feature exractor name. Our code is compatible with CLIP, DINO and DINO_v2. For more details, see configs/musc.yaml.
• --pretrained: Pretrained CLIP model. openai, laion400m_e31, and laion400m_e32 are optional.
• --feature_layers: The layers for extracting features in backbone(ViT).
• --img_resize: The size of the image inputted into the model.
• --divide_num: The number of subsets the whole test set is divided into.
• --r_list: The aggregation degrees of our LNAMD module.
• --output_dir: The directory that saves the anomaly prediction maps and metrics. This directory will be automatically created.
• --vis: Whether to save the anomaly prediction maps.
• --vis_type: Choose between single_norm and whole_norm. This means whether to normalize a single anomaly map or all of them together when visualizing.
• --save_excel: Whether to save anomaly classification and segmentation results (metrics).

💎Classification optimization (RsCIN): <a href='#all_catelogue'>[Back to Catalogue]</a>

We provide additional code in ./models/RsCIN_features folder to optimize the classification results of other methods using our RsCIN module. We use ViT-large-14-336 of CLIP to extract the image features of the MVTec AD and VisA datasets and store them in mvtec_ad_cls.dat and visa_cls.dat respectively. We show how to use them in ./models/RsCIN_features/RsCIN.py.

Example

Before using our RsCIN module, move RsCIN.py, mvtec_ad_cls.dat and visa_cls.dat to your project directory.

import numpy as np
from RsCIN import Mobile_RsCIN

classification_results = np.random.rand(83) # the classification results of your method.
dataset_name = 'mvtec_ad' # dataset name
class_name = 'bottle' # category name in the above dataset
optimized_classification_results = Mobile_RsCIN(classification_results, dataset_name=dataset_name, class_name=class_name)

The optimized_classification_results are the anomaly classification scores optimized by our RsCIN module.

Apply to the custom dataset

You can extract the image features of each image in the custom dataset, and store them in the variable cls_tokens. The multiple window sizes in the Multi-window Mask Operation can be adjusted by the value of k_list.

import numpy as np
from RsCIN import Mobile_RsCIN

classification_results = np.random.rand(83) # the classification results of your method.
cls_tokens = np.random.rand(83, 768)  # shape[N, C] the image features, N is the number of images
k_list = [2, 3] # the multiple window sizes in the Multi-window Mask Operation
optimized_classification_results = Mobile_RsCIN(classification_results, k_list=k_list, cls_tokens=cls_tokens)

🎖️Results of different datasets: <a href='#all_catelogue'>[Back to Catalogue]</a>

All the results are implemented by the default settings in our paper.

MVTec AD

VisA

BTAD

🎖️Results of different backbones: <a href='#all_catelogue'>[Back to Catalogue]</a>

The default backbone (feature extractor) in our paper is ViT-large-14-336 of CLIP. We also provide the supported codes for other image encoder of CLIP, DINO and DINO_v2. For more details, see configs/musc.yaml.

MVTec AD

VisA

BTAD

⌛Inference Time: <a href='#all_catelogue'>[Back to Catalogue]</a>

We show the inference time per image in the table below when using different backbones and image sizes. The default setting for number of images in mutual scoring module is 200, and GPU is NVIDIA RTX 3090.

🙋🙋‍♂️Frequently Asked Questions: <a href='#all_catelogue'>[Back to Catalogue]</a>

Q: Why do large areas of high anomaly scores appear on normal images in the visualization?

A: In the visualization, in order to highlight abnormal areas, we adopt a single anomaly map normalization by default. Even if the overall response of the single map is low, a large number of highlighted areas will appear after normalization. Normalization of all the anomaly maps together can be achieved by adding the vis_type parameter to the shell script and setting it as whole_norm, or by modifying the testing->vis_type parameter in the ./configs/musc.yaml.

Q: How to set the appropriate input image resolution ?

A: The image resolution img_resize input into the backbone is generally set to a multiple of the patch size of ViT. The commonly used values are 224, 240, 256, 336, 512 and 518. In the previous section <a href='#results_backbones'>(jump)</a>, we show the two input image resolutions commonly used by different feature extractors for reference. The image resolution can be changed by modifying the 'img_resize' parameter in the shell script, or by modifying the datasets->img_resize parameter in the ./configs/musc.yaml configuration file.

Citation: <a href='#all_catelogue'>[Back to Catalogue]</a>

@inproceedings{Li2024MuSc,
  title={MuSc: Zero-Shot Industrial Anomaly Classification and Segmentation with Mutual Scoring of the Unlabeled Images},
  author={Li, Xurui and Huang, Ziming and Xue, Feng and Zhou, Yu},
  booktitle={International Conference on Learning Representations},
  year={2024}
}

Thanks: <a href='#all_catelogue'>[Back to Catalogue]</a>

Our repo is built on PatchCore and APRIL-GAN, thanks their clear and elegant code !

License: <a href='#all_catelogue'>[Back to Catalogue]</a>

MuSc is released under the MIT Licence, and is fully open for academic research and also allow free commercial usage. To apply for a commercial license, please contact yuzhou@hust.edu.cn.