迈向工业异常检测的完全召回

本仓库包含 Roth 等人（2021）提出的 PatchCore 算法实现，详见 https://arxiv.org/abs/2106.08265。

仓库还提供了多种预训练模型，这些模型在图像级异常检测 AUROC 上可达 99.6%，在像素级异常定位 AUROC 上可达 98.4%，PRO 分数超过 95%（由于许可原因，未包含后者指标）。

defect_segmentation

如有问题和反馈，请联系 karsten.rh1@gmail.com！

快速指南

首先，克隆本仓库并设置 PYTHONPATH 环境变量，命令为 env PYTHONPATH=src python bin/run_patchcore.py。要在 MVTec AD 上训练 PatchCore（如下所述），请运行

datapath=/path_to_mvtec_folder/mvtec datasets=('bottle' 'cable' 'capsule' 'carpet' 'grid' 'hazelnut'
'leather' 'metal_nut' 'pill' 'screw' 'tile' 'toothbrush' 'transistor' 'wood' 'zipper')
dataset_flags=($(for dataset in "${datasets[@]}"; do echo '-d '$dataset; done))


python bin/run_patchcore.py --gpu 0 --seed 0 --save_patchcore_model \
--log_group IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0 --log_online --log_project MVTecAD_Results results \
patch_core -b wideresnet50 -le layer2 -le layer3 --faiss_on_gpu \
--pretrain_embed_dimension 1024  --target_embed_dimension 1024 --anomaly_scorer_num_nn 1 --patchsize 3 \
sampler -p 0.1 approx_greedy_coreset dataset --resize 256 --imagesize 224 "${dataset_flags[@]}" mvtec $datapath

该配置在 MVTec 224x224 尺寸图像上运行 PatchCore，使用在 ImageNet 上预训练的 WideResNet50 骨干网络。对于使用不同骨干网络、更大图像或集成模型的其他样本运行，请参见 sample_training.sh。

给定预训练的 PatchCore 模型（或所有 MVTec AD 子数据集的模型），可使用以下方式对其进行评估

datapath=/path_to_mvtec_folder/mvtec
loadpath=/path_to_pretrained_patchcores_models
modelfolder=IM224_WR50_L2-3_P001_D1024-1024_PS-3_AN-1_S0
savefolder=evaluated_results'/'$modelfolder

datasets=('bottle'  'cable'  'capsule'  'carpet'  'grid'  'hazelnut' 'leather'  'metal_nut'  'pill' 'screw' 'tile' 'toothbrush' 'transistor' 'wood' 'zipper')
dataset_flags=($(for dataset in "${datasets[@]}"; do echo '-d '$dataset; done))
model_flags=($(for dataset in "${datasets[@]}"; do echo '-p '$loadpath'/'$modelfolder'/models/mvtec_'$dataset; done))

python bin/load_and_evaluate_patchcore.py --gpu 0 --seed 0 $savefolder \
patch_core_loader "${model_flags[@]}" --faiss_on_gpu \
dataset --resize 366 --imagesize 320 "${dataset_flags[@]}" mvtec $datapath

此处托管了一组预训练的 PatchCore 模型：add link。若要使用这些模型（并复现训练过程），请查看 sample_evaluation.sh 和 sample_training.sh。

详细说明

环境要求

我们的结果是使用 Python 3.8 计算得出的，相关包及其版本信息已在 requirements.txt 中注明。通常情况下，大多数实验的 GPU 内存占用不应超过 11GB；但使用过大的输入图像会导致更高的内存消耗。

MVTec AD 数据集的设置

如需设置主要的 MVTec AD 基准数据集，请从此处下载：https://www.mvtec.com/company/research/datasets/mvtec-ad。将其放置在某个路径 datapath 下，并确保其遵循以下数据目录结构：

mvtec
|-- bottle
|-----|----- ground_truth
|-----|----- test
|-----|--------|------ good
|-----|--------|------ broken_large
|-----|--------|------ ...
|-----|----- train
|-----|--------|------ good
|-- cable
|-- ...

共包含 15 个子数据集：bottle、cable、capsule、carpet、grid、hazelnut、leather、metal_nut、pill、screw、tile、toothbrush、transistor、wood、zipper。

“训练” PatchCore

PatchCore 提取预训练的、经过局部聚合的训练补丁特征的（核心集子采样）记忆库：

patchcore_architecture

为此，我们提供了 bin/run_patchcore.py，它使用 click 来管理和聚合输入参数。其使用方式大致如下

python bin/run_patchcore.py \
--gpu <gpu_id> --seed <seed> # Set GPU-id & reproducibility seed.
--save_patchcore_model # If set, saves the patchcore model(s).
--log_online # If set, logs results to a Weights & Biases account.
--log_group IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0 --log_project MVTecAD_Results results # Logging details: Name of the run & Name of the overall project folder.

patch_core  # We now pass all PatchCore-related parameters.
-b wideresnet50  # Which backbone to use.
-le layer2 -le layer3 # Which layers to extract features from.
--faiss_on_gpu # If similarity-searches should be performed on GPU.
--pretrain_embed_dimension 1024  --target_embed_dimension 1024 # Dimensionality of features extracted from backbone layer(s) and final aggregated PatchCore Dimensionality
--anomaly_scorer_num_nn 1 --patchsize 3 # Num. nearest neighbours to use for anomaly detection & neighbourhoodsize for local aggregation.

sampler # We now pass all the (Coreset-)subsampling parameters.
-p 0.1 approx_greedy_coreset # Subsampling percentage & exact subsampling method.

dataset # We now pass all the Dataset-relevant parameters.
--resize 256 --imagesize 224 "${dataset_flags[@]}" mvtec $datapath # Initial resizing shape and final imagesize (centercropped) as well as the MVTec subdatasets to use.

请注意，sample_runs.sh 包含用于实现强大 AD 性能的示例性训练运行。由于代码库变更（以及硬件差异），结果可能与论文中报告的结果略有偏差，但总体上应非常接近甚至更好。如前所述，为了便于重复使用和复现，我们还提供了多个预训练的 PatchCore 模型，托管于 add link - 下载文件夹，解压后，将所选模型传递给 bin/load_and_evaluate_patchcore.py，该脚本展示了示例性的评估过程。

在训练期间（之后），将存储以下信息：

|PatchCore model (if --save_patchcore_model is set)
|-- models
|-----|----- mvtec_bottle
|-----|-----------|------- nnscorer_search_index.faiss
|-----|-----------|------- patchcore_params.pkl
|-----|----- mvtec_cable
|-----|----- ...
|-- results.csv # Contains performance for each subdataset.

|Sample_segmentations (if --save_segmentation_images is set)

除主要训练流程外，我们还集成了Weights-&-Biases日志功能，该功能允许您将所有训练和测试性能在线记录到Weights-and-Biases服务器（https://wandb.ai）。要使用此功能，请添加--log_online标志，并在run_patchcore.py > --log_wandb_key中提供您的W&B密钥。

最后，鉴于PatchCore的有效性和高效性，我们还加入了使用骨干网络和网络特征图集成的选项。为此，请使用-b <backbone>提供要使用的骨干网络列表（如/src/anomaly_detection/backbones.py中所列），并根据它们的顺序，使用-le idx.<layer_name>指定要提取的层。使用三个不同骨干网络的示例如下所示

python bin/run_patchcore.py --gpu <gpu_id> --seed <seed> --save_patchcore_model --log_group <log_name> --log_online --log_project <log_project> results \

patch_core -b wideresnet101 -b resnext101 -b densenet201 -le 0.layer2 -le 0.layer3 -le 1.layer2 -le 1.layer3 -le 2.features.denseblock2 -le 2.features.denseblock3 --faiss_on_gpu \

--pretrain_embed_dimension 1024  --target_embed_dimension 384 --anomaly_scorer_num_nn 1 --patchsize 3 sampler -p 0.01 approx_greedy_coreset dataset --resize 256 --imagesize 224 "${dataset_flags[@]}" mvtec $datapath

当使用 --save_patchcore_model 时，在集成情况下，会存储相应的 PatchCore 参数集成。

评估预训练的 PatchCore 模型

要评估一个/我们的预训练 PatchCore 模型，运行

python bin/load_and_evaluate_patchcore.py --gpu <gpu_id> --seed <seed> $savefolder \
patch_core_loader "${model_flags[@]}" --faiss_on_gpu \
dataset --resize 366 --imagesize 320 "${dataset_flags[@]}" mvtec $datapath

假设您的预训练模型位置包含在 model_flags 中；每个子数据集对应 dataset_flags 中的一项。结果将存储在 savefolder 中。模型和数据集标志示例：

model_flags=('-p', 'path_to_mvtec_bottle_patchcore_model', '-p', 'path_to_mvtec_cable_patchcore_model', ...)
dataset_flags=('-d', 'bottle', '-d', 'cable', ...)

预训练模型的预期性能

尽管由于软硬件差异可能导致性能略有变化，但所提供的预训练模型应能达到其各自 results.csv 文件中记录的性能。平均性能（特别是基线模型 WR50 以及更大的集成模型）大致如下：

模型	平均 AUROC	平均分割 AUROC	平均 PRO
WR50-baseline	99.2%	98.1%	94.4%
Ensemble	99.6%	98.2%	94.9%

引用

如果您使用本仓库中的代码，请引用

@misc{roth2021total,
      title={Towards Total Recall in Industrial Anomaly Detection},
      author={Karsten Roth and Latha Pemula and Joaquin Zepeda and Bernhard Schölkopf and Thomas Brox and Peter Gehler},
      year={2021},
      eprint={2106.08265},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

安全

有关更多信息，请参见 CONTRIBUTING。

许可证

本项目基于 Apache-2.0 许可证授权。