MMR: A Python repository from ZhaoZhibin

Industrial Anomaly Detection with Domain Shift: A Real-world Dataset and Masked Multi-scale Reconstruction

This is an official PyTorch implementation of the paper Industrial Anomaly Detection with Domain Shift: A Real-world Dataset and Masked Multi-scale Reconstruction.

@article{zhang2023industrial,
  title={Industrial Anomaly Detection with Domain Shift: A Real-world Dataset and Masked Multi-scale Reconstruction},
  author={Zhang, Zilong and Zhao, Zhibin and Zhang, Xingwu and Sun, Chuang and Chen, Xuefeng},
  journal={arXiv preprint arXiv:2304.02216},
  year={2023}
}

Datasets

We release a real-world Aero-engine Blade Anomaly Detection (AeBAD) dataset, consisting of two sub-datasets: the single-blade dataset (AeBAD-S) and the video anomaly detection dataset of blades (AeBAD-V). Compared to existing datasets, AeBAD has the following two characteristics: 1.) The target samples are not aligned and at different sacles. 2.) There is a domain shift between the distribution of normal samples in the test set and the training set, where the domain shifts are mainly caused by the changes in illumination and view.

Download dataset at here (Google Drive) or here (access code: g4pr) (Tian Yi Yun Pan).

AeBAD-S

AeBAD-V

Visualization for Videos

①: Original Video ②: PatchCore ③: ReverseDistillation ④: DRAEM ⑤: NSA ⑥: MMR

video 1

①

②

③

④

⑤

⑥

Video 2

①

②

③

④

⑤

⑥

Video 3

①

②

③

④

⑤

⑥

Get Started

Pre-trained models

Download the pre-trained model of MAE (ViT-base) at here.

Dataset

MVTec:

Create the MVTec dataset directory. Download the MVTec-AD dataset from here. The MVTec dataset directory should be as follows.

|-- data
    |-- MVTec-AD
        |-- mvtec_anomaly_detection
            |-- object (bottle, etc.)
                |-- train
                |-- test
                |-- ground_truth

AeBAD:

Download the AeBAD dataset from the above link. The AeBAD dataset directory should be as follows.

|-- AeBAD
    |-- AeBAD_S
        |-- train
            |-- good
                |-- background
        |-- test
                |-- ablation
                    |-- background
        |-- ground_truth
                |-- ablation
                    |-- view
    |-- AeBAD_V
        |-- test
            |-- video1
                |-- anomaly
        |-- train
            |-- good
                |-- video1_train

Note that background, view and illumination in the train set is different from test. The background, view and illumination in test is unseen for the training set.

Virtual Environment

Use the following commands:

pip install -r requirements.txt

Train and Test for MVTec, AeBAD

Train the model and evaluate it for each category or different domains. This will output the results (sample-level AUROC, pixel-level AUROC and PRO) for each category. It will generate the visualization in the directory.

run the following code:

sh mvtec_run.sh

sh AeBAD_S_run.sh

sh AeBAD_V_run.sh

TRAIN.MMR.model_chkpt in MMR.yaml is the path of above download model. TRAIN.dataset_path (TEST.dataset_path) is the path of data. Set Test.save_segmentation_images as True or False to save processed image.

Note that for AeBAD-V, we only evaluate the sample-level metric. The pixel-level metric is 0.

ZhaoZhibin/MMR