基于OpenVINO实现无监督异常检测

Color Space

发布于 2024-06-04 13:32:54

4870

发布于 2024-06-04 13:32:54

视觉/图像重磅干货，第一时间送达!

异常检测(AD) 在欺诈检测、网络安全和医疗诊断等关键任务应用中至关重要。由于数据的高维性和底层模式的复杂性，图像、视频和卫星图像等视觉数据中的异常检测尤其具有挑战性。然而，视觉异常检测对于检测制造中的缺陷、识别监控录像中的可疑活动以及检测医学图像中的异常至关重要。

在本文中，您将学习如何使用OpenVINO 工具包中的FiftyOne和Anomalib对视觉数据执行异常检测。为了演示，我们将使用MVTec AD 数据集，其中包含具有划痕、凹痕和孔洞等异常的各种物体的图像。

它涵盖以下内容：

在 FiftyOne 中加载 MVTec AD 数据集
使用 Anomalib 训练异常检测模型
评估 FiftyOne 中的异常检测模型

安装依赖项

确保你在虚拟环境中运行它python=3.10。Anomalib 需要 Python 3.10，因此请确保你安装了正确的版本。

conda create -n anomalib_env python=3.10
conda activate anomalib_env

此后，按照Anomalib README中的说明从源代码安装 Anomalib及其依赖项。这些在 Google Colab 上可能需要一些时间，但对于本地安装应该很快：

pip install -U torchvision einops FrEIA timm open_clip_torch imgaug lightning kornia openvino git+https://github.com/openvinotoolkit/anomalib.git

我们准备安装更多软件包。现在您可以明白为什么我们建议为该项目使用虚拟环境！

huggingface_hub用于加载 MVTec AD 数据集
clip用于计算图像嵌入
umap-learn用于降维

pip install -U huggingface_hub umap-learn git+https://github.com/openai/CLIP.git

加载和可视化 MVTec AD 数据集

现在，让我们从FiftyOne导入我们需要的所有相关模块：

import fiftyone as fo # 基础库和应用程序
import fiftyone.brain as fob # ML 方法
import fiftyone.zoo as foz # zoo 数据集和模型
from fiftyone import ViewField as F # 定义视图的助手
import fiftyone.utils.huggingface as fouh # Hugging Face 集成

并从 Hugging Face Hub 加载 MVTec AD 数据集：

dataset = fouh.load_from_hub("Voxel51/mvtec-ad", persistent=True, overwrite=True)

在继续之前，让我们看一下FiftyOne 应用程序中的数据集：

session = fo.launch_app(dataset)

该数据集包含 12 个对象类别的 5354 张图像。每个类别都有“良好”和“异常”图像，这些图像存在划痕、凹痕和孔洞等缺陷。每个异常样本还带有一个掩模，用于定位图像中的缺陷区域。

缺陷标签因类别而异，这在现实世界的异常检测场景中很常见。在这些场景中，您会为每个类别训练不同的模型。在这里，我们将介绍一个类别的流程，您可以将相同的步骤应用于其他类别。

还有一点需要注意的是，数据集被分为训练集和测试集。训练集只包含“良好”图像，而测试集则包含“良好”和“异常”图像。

在训练模型之前，让我们深入研究数据集。通过计算图像嵌入并在低维空间中可视化它们，我们可以了解数据中隐藏的结构和模式。首先，我们将使用 CLIP 模型计算数据集中所有图像的嵌入：

model = foz.load_zoo_model( "clip-vit-base32-torch" )   # 从 zoo 加载 CLIP 模型

# 计算数据集的嵌入
dataset.compute_embeddings( 
    model=model, embeddings_field= "clip_embeddings" , batch_size= 64
 ) 

# 使用 UMAP 对嵌入进行降维
fob.compute_visualization( 
    dataset, embeddings= "clip_embeddings" , method= "umap" , brain_key= "clip_vis"
 )

刷新 FiftyOne 应用程序，单击“+”选项卡，然后选择“Embeddings”。从下拉菜单中选择“all_clip_vis”。您将看到 2D 空间中图像嵌入的散点图，其中每个点对应于数据集中的一个样本。

使用颜色下拉菜单，注意嵌入如何根据对象类别进行聚类。这是因为 CLIP 对图像的语义信息进行编码。此外，CLIP 嵌入不会根据缺陷类型在类别内进行聚类。

训练异常检测模型

现在我们对数据集有了了解，我们准备使用 Anomalib 训练异常检测模型。

任务：Anomalib 支持图像的分类、检测和分割任务。我们将重点关注分割，其中模型预测图像中的每个像素是否异常，并创建一个定位缺陷的掩码。

模型：Anomalib 支持多种异常检测算法。在本演练中，我们将使用两种算法：

PaDiM：用于异常检测和定位的补丁分布建模框架
PatchCore：迈向工业异常检测的全面召回

预处理：在训练模型之前，我们将在本演练中将图像大小调整为 256x256 像素。通过 Torchvision 的 Resize 类将其添加为转换，我们可以在训练和推理过程中动态调整图像大小。

从Anomalib 和辅助模块导入处理图像和路径所需的模块：

import numpy as np
import os
from pathlib import Path
from PIL import Image
from torchvision.transforms.v2 import Resize

from anomalib import TaskType
from anomalib.data.image.folder import Folder
from anomalib.deploy import ExportType, OpenVINOInferencer
from anomalib.engine import Engine
from anomalib.models import Padim, Patchcore

现在，定义一些在整个笔记本中使用的常量。

OBJECT：我们将重点关注的对象类别。在本演练中，我们将使用“瓶子”。如果您想要循环遍历类别，可以使用 dataset.distinct("category.label") 从数据集中获取类别列表。
ROOT_DIR：Anomalib 将在其中查找图像和掩码的根目录。我们的数据已存储在磁盘上，因此我们只需将文件符号链接到 Anomalib 所需的目录即可。
TASK：我们正在执行的任务。我们将在本演练中使用“分段”。
IMAGE_SIZE：在训练模型之前调整图像的大小。我们将使用 256x 256 像素。

OBJECT = "bottle"  ## 要训练的对象
ROOT_DIR = Path( "/tmp/mvtec_ad" ) ## 用于存储 anomalib 数据的根目录
TASK = TaskType.SEGMENTATION ## 模型的任务类型
IMAGE_SIZE = ( 256 , 256 ) ## 预处理图像大小以保证均匀性

对于给定的对象类型（类别），create_datamodule()下面的函数会创建一个 AnomalibDataModule对象。这将被传递到我们引擎的fit()方法来训练模型，并用于实例化数据加载器以进行训练和验证。

代码可能看起来很复杂，所以让我们分解一下发生了什么：

我们创建的数据子集仅包含“良好”的训练图像和“异常”图像以供验证。
我们将图像和掩码符号链接到 Anomalib 期望的目录。
我们从 Anomalib 实例化并设置一个数据模块Folder，它是自定义数据集的通用类。

💡 也可以DataLoader从头开始创建一个 torch 并将其传递给引擎的fit()方法。这可以让你更好地控制数据加载过程。这留给读者练习 😉。

def create_datamodule(object_type, transform=None):
    ## Build transform
    if transform is None:
        transform = Resize(IMAGE_SIZE, antialias=True)

    normal_data = dataset.match(F("category.label") == object_type).match(
        F("split") == "train"
    )
    abnormal_data = (
        dataset.match(F("category.label") == object_type)
        .match(F("split") == "test")
        .match(F("defect.label") != "good")
    )

    normal_dir = Path(ROOT_DIR) / object_type / "normal"
    abnormal_dir = ROOT_DIR / object_type / "abnormal"
    mask_dir = ROOT_DIR / object_type / "mask"

    # create directories if they do not exist
    os.makedirs(normal_dir, exist_ok=True)
    os.makedirs(abnormal_dir, exist_ok=True)
    os.makedirs(mask_dir, exist_ok=True)

    if not os.path.exists(str(normal_dir)):
        normal_data.export(
            export_dir=str(normal_dir),
            dataset_type=fo.types.ImageDirectory,
            export_media="symlink",
        )

    for sample in abnormal_data.iter_samples():
        base_filename = sample.filename
        dir_name = os.path.dirname(sample.filepath).split("/")[-1]
        new_filename = f"{dir_name}_{base_filename}"
        if not os.path.exists(str(abnormal_dir / new_filename)):
            ## symlink anomalous image into Anomalib abnormal dir
            os.symlink(sample.filepath, str(abnormal_dir / new_filename))


        if not os.path.exists(str(mask_dir / new_filename)):
            ## symlink mask into Anomalib mask dir
            os.symlink(sample.defect_mask.mask_path, str(mask_dir / new_filename))


    ## Create a DataModule in Anomalib
    datamodule = Folder(
        name=object_type,
        root=ROOT_DIR,
        normal_dir=normal_dir,
        abnormal_dir=abnormal_dir,
        mask_dir=mask_dir,
        task=TASK,
        transform=transform
    )
    datamodule.setup()
    return datamodule

现在，我们可以将所有内容整合在一起。train_and_export_model()下面的函数使用 Anomalib 的类训练异常检测模型Engine，将模型导出到 OpenVINO，并返回模型“推理器”对象。推理器对象用于对新图像进行预测。

def  train_and_export_model ( object_type, model, transform= None ): 
    ## 在我们的数据上训练模型
    datamodule = create_datamodule(object_type, transform=transform) 
    engine = Engine(task=TASK) 
    engine.fit(model=model, datamodule=datamodule) 

    ## 将模型导出为 OpenVINO 格式以进行快速推理
    engine.export( 
        model=model, 
        export_type=ExportType.OPENVINO, 
    ) 
    output_path = Path(engine.trainer.default_root_dir) 

    openvino_model_path = output_path / "weights" / "openvino" / "model.bin"
     metadata = output_path / "weights" / "openvino" / "metadata.json" 


    ## 从导出加载推理对象 inferencer
     = OpenVINOInferencer( 
        path=openvino_model_path, 
        metadata=metadata, 
        device= "CPU" , 
    ) 
    return inferencer

我们先尝试PaDiM一下。训练过程应该不到一分钟：

model = Padim()

inferencer = train_and_export_model(OBJECT, model)

就这样，我们就有了一个针对“瓶子”类别进行训练的异常检测模型。让我们在单个图像上运行推理器并检查结果：

## get the test split of the dataset
test_split = dataset.match(F("category.label") == OBJECT).match(F("split") == "test")

## get the first sample from the test split
test_image = Image.open(test_split.first().filepath)

output = inferencer.predict(image=test_image)
print(output)

ImageResult(image=[[[255 255 255]
  [255 255 255]
  [255 255 255]
  ...
  [255 255 255]
  [255 255 255]
  [255 255 255]]

  ...
  [255 255 255]
  [255 255 255]
  [255 255 255]]], pred_score=0.7751642969087686, pred_label=1, anomaly_map=[[0.32784402 0.32784402 0.32784414 ... 0.3314721  0.33147204 0.33147204]
 [0.32784402 0.32784402 0.32784414 ... 0.3314721  0.33147204 0.33147204]
 [0.32784408 0.32784408 0.3278442  ... 0.33147222 0.33147216 0.33147216]
 ...
 [0.32959    0.32959    0.32959005 ... 0.3336093  0.3336093  0.3336093 ]
 [0.3295899  0.3295899  0.32958996 ... 0.33360928 0.33360928 0.33360928]
 [0.3295899  0.3295899  0.32958996 ... 0.33360928 0.33360928 0.33360928]], gt_mask=None, gt_boxes=None, pred_boxes=None, box_labels=None, pred_mask=[[0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 ...
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]], heat_map=[[[153 235 255]
  [153 235 255]
  [153 235 255]
  ...
  [153 236 255]
  [153 236 255]
  [153 236 255]]
  ...
  [153 238 255]
  [153 238 255]
  [153 238 255]]], segmentations=[[[255 255 255]
  [255 255 255]
  [255 255 255]
  ...
  [255 255 255]
  [255 255 255]
  [255 255 255]]
  ...
  [255 255 255]
  [255 255 255]
  [255 255 255]]])

输出包含一个标量异常分数pred_score、一个pred_mask表示预测异常区域的和一个显示每个像素异常分数的热图 anomaly_map。这些都是理解模型预测的宝贵信息。

下面的函数run_inference()将以 FiftyOne 样本集合（例如我们的测试集）作为输入，以及推理器对象和用于将结果存储在样本中的键。它将对集合中的每个样本运行模型并存储结果。阈值参数充当异常分数的截止值。如果分数高于阈值，则样本被视为异常。在此示例中，我们将使用 0.5 的阈值，但您可以尝试使用不同的值。

def run_inference(sample_collection, inferencer, key, threshold=0.5):
    for sample in sample_collection.iter_samples(autosave=True, progress=True):
        output = inferencer.predict(image=Image.open(sample.filepath))

        conf = output.pred_score
        anomaly = "normal" if conf < threshold else "anomaly"

        sample[f"pred_anomaly_score_{key}"] = conf
        sample[f"pred_anomaly_{key}"] = fo.Classification(label=anomaly)
        sample[f"pred_anomaly_map_{key}"] = fo.Heatmap(map=output.anomaly_map)
        sample[f"pred_defect_mask_{key}"] = fo.Segmentation(mask=output.pred_mask)

让我们对测试分割进行推理，并在 FiftyOne 应用程序中可视化结果：

run_inference(test_split, inferencer, "padim")
session = fo.launch_app(view=test_split)

评估异常检测模型

我们有一个异常检测模型，但我们如何知道它是否好用？首先，我们可以使用精度、召回率和 F1 分数指标来评估模型。FiftyOne 的评估 API使这变得简单。我们将评估模型的全图像分类性能以及分割性能。

我们需要准备评估数据。首先，我们需要为“正常”图像添加空掩码，以确保评估公平：

for sample in test_split.iter_samples(autosave=True, progress=True):
    if sample["defect"].label == "good":
        sample["defect_mask"] = fo.Segmentation(
            mask=np.zeros_like(sample["pred_defect_mask_padim"].mask)
        )

我们还需要确保真实值和预测值之间的命名/标签的一致性。我们将所有“良好”图像重命名为“正常”，并将每种类型的异常重命名为“异常”：

old_labels = test_split.distinct("defect.label")
label_map = {label:"anomaly" for label in old_labels if label != "good"}
label_map["good"] = "normal"
mapped_view = test_split.map_labels("defect", label_map)
session.view = mapped_view.view()

对于分类，我们将使用二元评估，其中“正常”为负类，“异常”为正类：

eval_classif_padim = mapped_view.evaluate_classifications(
    "pred_anomaly_padim",
    gt_field="defect",
    eval_key="eval_classif_padim",
    method="binary",
    classes=["normal", "anomaly"],
)
eval_classif_padim.print_report()

               precision    recall  f1-score   support

      normal       0.95      0.90      0.92        20
     anomaly       0.97      0.98      0.98        63

    accuracy                           0.96        83
   macro avg       0.96      0.94      0.95        83
weighted avg       0.96      0.96      0.96        83

比较异常检测模型

虽然异常检测是无监督的，但这并不意味着我们不能比较模型并选择最适合我们用例的模型。我们可以在同一数据上训练多个模型，并使用 F1 分数、准确率和召回率等指标比较它们的性能。我们还可以通过检查它们生成的掩码和热图来直观地比较模型。

我们重复一下PatchCore模型的训练过程，并比较一下这两个模型：

## 训练 Patchcore 模型并运行推理

model = Patchcore() 

## 这将需要更长的时间来训练，但仍应少于 5 分钟
inferencer = train_and_export_model(OBJECT, model) 

run_inference(mapped_view, inferencer, "patchcore" ) 

## 在分类任务上评估 Patchcore 模型
eval_classif_patchcore = tagged_view.evaluate_classifications( 
    "pred_anomaly_patchcore" , 
    gt_field= "defect" , 
    eval_key= "eval_classif_patchcore" , 
    method= "binary" , 
    classes=[ "normal" , "anomaly" ], 
) 

eval_classif_patchcore.print_report()

              precision    recall  f1-score   support

      normal       0.95      1.00      0.98        20
     anomaly       1.00      0.98      0.99        63

    accuracy                           0.99        83
   macro avg       0.98      0.99      0.98        83
weighted avg       0.99      0.99      0.99        83

eval_seg_patchcore = mapped_view.match(F("defect.label") == "anomaly").evaluate_segmentations(
    "pred_defect_mask_patchcore",
    gt_field="defect_mask",
    eval_key="eval_seg_patchcore",
)
eval_seg_patchcore.print_report(classes=[0, 255])
session.view = mapped_view.shuffle().view()

      precision    recall  f1-score   support

           0       0.99      0.95      0.97 47143269.0
         255       0.60      0.85      0.70 3886731.0

   micro avg       0.95      0.95      0.95 51030000.0
   macro avg       0.80      0.90      0.84 51030000.0
weighted avg       0.96      0.95      0.95 51030000.0

这些指标支持了我们在应用程序中看到的结果：PatchCore 对“异常”类别的召回率更高，但准确率较低。这意味着它更有可能发现异常，但也更有可能做出误报预测。毕竟，PatchCore 是为工业异常检测中的“全面召回”而设计的。