PatchCore的代码复现

最新推荐文章于 2025-11-18 10:13:20 发布

原创最新推荐文章于 2025-11-18 10:13:20 发布 · 2.2k 阅读

CC 4.0 BY-SA版权

文章标签：

前言
在这里插入图片描述
源码地址：https://github.com/amazon-science/patchcore-inspection
论文地址：https://arxiv.org/abs/2106.08265
异常检测问题在工业图像数据分析中扮演着至关重要的角色，其目的是从大量正常数据中识别出异常行为或模式。这一任务的挑战在于，正常数据的样本相对容易获取，而异常情况却因其稀有性和多样性而难以收集。为了解决这一问题，研究者们开发了多种方法，其中一种备受关注的方法是PatchCore模型。
PatchCore模型是一种先进的工业异常检测方法，它在MVTec数据集上取得了最先进的性能（State of the Art，简称SOTA。这个数据集是工业领域内公认的用于评估异常检测算法的标准数据集，包含了多种不同类别的工业产品图像，既有正常样本也有异常样本。
PatchCore的关键优势在于它的特征提取机制。传统的异常检测方法通常需要大量的标注数据来训练模型，以学习正常和异常图像之间的差异。然而，PatchCore采用了一种不同的策略，它利用了预训练模型（如WideResNet50）来提取图像特征。这种方法的优势在于，预训练模型已经在大型数据集（如ImageNet）上学习了丰富的视觉特征，因此可以直接应用于新的任务，而无需进行额外的特征提取训练。
这里只复现代码中训练和测试流程，不对论文核心进行分析。
使用数据
数据来源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（拉链）。

模型训练

run_patchcore.py
--gpu 0 # 设置GPU ID
--seed 0 # 重复性种子 
--save_patchcore_model # 如果设置，则保存PatchCore模型
--save_segmentation_images
--log_group IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0
--log_project MVTecAD_Results results # 日志详细信息：运行名称和整体项目文件夹的名称
patch_core # 现在传递所有与PatchCore相关的参数
-b wideresnet50  # 使用哪个基础网络
-le layer2 # 从哪些层提取特征
-le layer3
--faiss_on_gpu # 是否在GPU上执行相似性搜索
--pretrain_embed_dimension 1024
--target_embed_dimension 1024 # 从基础层提取的特征维度和最终聚合的PatchCore维度
--anomaly_scorer_num_nn 1 # 异常检测使用的最近邻数量
--patchsize 3 # 和局部聚合的邻居大小
sampler # 现在传递所有（Coreset-）子采样的参数
-p 0.1 approx_greedy_coreset # # 子采样百分比和精确的子采样方法
dataset # 现在传递所有与数据集相关联的参数
--resize 256 # 初始缩放形状
--imagesize 224 # 最终裁剪图片尺寸
-d wood # 选择的数据类型
mvtec \datasets\mvtec # 要使用的MVTec子数据集

出现问题1：保存结果图时总是会将原图覆盖掉，检查后发现src/patchcore/utils.py代码存在bug，进行了修改

savename = image_path.split("/")
savename = "_".join(savename[-save_depth:])
savename = os.path.join(savefolder, savename)
这三行删掉，修改成
folder_name = os.path.basename(os.path.dirname(image_path))
save_path = os.path.join(savefolder, folder_name)
if not os.path.exists(save_path):
	os.makedirs(save_path)
savename = os.path.join(save_path, os.path.basename(image_path))

出现问题2：AttributeError: ‘MVTecDataset’ object has no attribute ‘transform_std’
因为主干是基于imagenet数据集训练的，所以均值和方差都是imagenet，因此可以在patchcore/datasets/mvtec.py中，在类class MVTecDataset 的def init 的函数中加入下面两行

self.transform_std = IMAGENET_STD
self.transform_mean = IMAGENET_MEAN

模型推理

load_and_evaluate_patchcore.py
--gpu 0
--seed 0
--save_segmentation_images
"results/MVTecAD_Results/IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0/evaluateAnswer" # 保存测试结果的路径
patch_core_loader
--patch_core_paths
"results/MVTecAD_Results/IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0/models/mvtec_tile" # 前训练的模型的路径
--faiss_on_gpu
dataset
--resize 256
--imagesize 224
--num_workers 0
-d tile # 选择的数据类型
mvtec \datasets\mvtec # 要使用的MVTec子数据集