工业级KNN异常检测：基于torch和timm的实现

葛瀚纲Deirdre

于 2024-08-31 09:45:06 发布

阅读量550

点赞数 5

本文链接：https://blog.csdn.net/gitblog_00108/article/details/141745978

版权

工业级KNN异常检测：基于torch和timm的实现

ind_knn_adIndustrial knn-based anomaly detection for images. Visit streamlit link to check out the demo.项目地址:https://gitcode.com/gh_mirrors/in/ind_knn_ad

项目介绍

本项目工业级KNN异常检测(ind_knn_ad) 是一个用PyTorch和timm框架实现的图像异常检测工具包。它主要关注于通过K近邻(KNN)方法在不同特征空间中进行异常检测。该仓库复现了包括SPADEFootnote:[Cohen et al 2021], PaDiMFootnote:[Defard et al 2020]（注释：虽然PaDiM不直接使用KNN，但在实现层面有许多共通点）以及PatchCoreFootnote:[Roth et al 2021]等在内的几种KNN基础的异常检测算法。它特别适用于工业视觉场景，如缺陷检测，利用预训练模型和自定义数据集来识别正常与异常图像。

项目快速启动

首先，确保你的环境中已安装必要的Python环境管理工具如pipenv或conda，并且已配置好Python 3.8+。以下步骤将引导你快速开始：

安装依赖

使用pipenv进行依赖管理，确保高效的虚拟环境创建与隔离：

pipenv install --dev

这将会根据requirements-dev.txt文件安装所有必需的开发与运行时依赖。

运行示例

选择一个方法（比如SPADECite:[Cohen et al 2021]），并准备或选择MVTec AD数据集的一个类别作为训练数据：

python run.py SPADE --dataset bottle

此命令会训练一个SPADECite:[Cohen et al 2021]模型，其中--dataset bottle可以替换为你想用于训练的数据集类别名。模型训练完成后，可以在/results/目录下找到评估结果。

代码示例

下面是如何初始化一个SPADECite:[Cohen et al 2021]模型并对其进行简单使用的代码段：

from indad.model import SPADE
model = SPADE(k=5, backbone_name="resnet18")
class_name = "bottle"
train_ds, test_ds = MVTecDataset(class_name).get_dataloaders()
model.fit(train_ds)
image_rocauc, pixel_rocauc = model.evaluate(test_ds)
print(f"Image ROC-AUC: {image_rocauc}, Pixel ROC-AUC: {pixel_rocauc}")