【2023工业3D异常检测文献】EasyNet: 基于多尺度多模态特征编码器-解码器的异常检测方法

qq_1532145264

于 2024-10-22 09:11:31 发布

阅读量1.1k

点赞数 15

分类专栏：机器学习文章标签： 3d 深度学习计算机视觉机器学习神经网络视觉检测图像处理

本文链接：https://blog.csdn.net/qq_1532145264/article/details/142764646

版权

机器学习专栏收录该内容

61 篇文章

订阅专栏

EasyNet: An Easy Network for 3D Industrial Anomaly Detection

1、Background

目前3D异常检测方法存在以下几个缺点：

基于需要特征提取的一些异常检测方法依赖于大型预训练模型的表征能力，导致推理速度慢和巨大的存储开销。
基于特征嵌入的3D-AD方法过度使用记忆库，在实际应用中导致巨大的记忆库成本。

为了避免使用大型预训练模型和记忆库，提出了一个简单但有效的多模态异常检测和定位网络EasyNet。

具体来说，EasyNet由两部分组成，多模态重建网络（MRN）和多模态分割网络（MSN）。

首先，不直接使用预训练特征，而是在RGB图像和深度图像上生成合成异常，用语义上合理的无内容重建原始图像，并获取多尺度特征。

同时，为了简化基于记忆库的异常检测过程，将异常和重建的多尺度特征输入到一个简单的MSN中以获得异常图。如图3所示，整个架构，包括MRN和MSN，显著促进了RGB和深度特征之间的交互。

为了减少RGB和深度图像之间的干扰，提出了一个基于注意力的信息熵融合模块。此外，发现一些3D-AD方法，如AST和BTF ，无法充分利用多模态融合的优势，即RGB-D性能并不比RGB性能有竞争力。主要原因是RGB或深度图像中没有统一的异常模式。例如，一些异常可以通过纯RGB图像检测到，而深度信息则作为噪声，可能会降低整体异常检测性能。

因此，提出了一个动态的多模态融合方案，以利用RGB和深度特征。融合方案的架构。

此外，提出的基于注意力的信息熵融合模块易于训练和应用，具有出色的性能和推理速度。

2、Method

在这里插入图片描述

EasyNet的总体架构主要由三部分组成：

异常生成器：
- 异常生成器的职责是向原始的多模态图像（RGB和深度图像）添加柏林噪声（Berlin noise）。
- 柏林噪声是一种随机噪声，用于模拟合成的异常图像和正常图像。
- 通过这种方式，异常生成器创建了用于训练网络的带噪声的图像样本。
多模态重建网络（MRN）：
- MRN负责执行图像重建任务，将增强后的合成异常图像恢复成无异常的RGB和深度图像。
- 在重建过程中，MRN从多个层次提取多模态特征信息，文中提到以两个层次作为示例。
- 这意味着MRN不仅重建图像，还从不同层次学习并提取特征，以便更好地理解图像内容。
多模态分割网络（MSN）：
- MSN的作用是融合提取的多模态特征，以产生精确的异常分割图。
- 它使用了一个基于注意力的信息熵融合模块来控制特征的融合。
- 在训练期间，该融合模块是完全激活的，而在推理（测试）期间，则通过计算自注意力信息熵分数来控制不同模式下特征信息的流动。

对于重建任务，EasyNet使用了：

SSIM损失（Structural Similarity Index Measure loss）和MSE损失（Mean Squared
Error loss）来计算RGB图像的重建损失。
仅使用MSE损失来计算深度图像的重建损失。
Focal损失用于计算像素分类任务的损失，这有助于解决类别不平衡问题，提高小样本类别的检测精度。

基于注意力的信息熵融合模块 。多模态分割网络（MSN）中，用于融合RGB和深度图像的特征。基于注意力的信息熵融合模块的核心思想是动态地选择最有信息量的特征进行融合，以提高异常检测的准确性。

在这里插入图片描述

基于注意力的信息熵融合模块方法：

特征输入：RGB特征和深度特征被送入融合模块。
注意力模块：利用注意力机制为RGB和深度特征分别生成注意力得分。注意力得分反映了不同通道特征的重要性。
信息熵计算：计算融合特征的信息熵，信息熵是衡量特征不确定性的指标。如果融合特征的信息熵高于RGB特征的信息熵，这意味着融合特征包含了更多的信息。
融合决策：通过比较融合特征的信息熵与RGB特征的信息熵，模块决定是否将深度特征与RGB特征进行融合。如果融合特征的信息熵高于RGB特征的信息熵加上一个阈值调整因子（α），则认为融合特征对异常检测更有帮助，并将它们融合。否则，只使用RGB特征进行异常检测。
输出融合特征：根据上述决策，输出融合后的特征或仅RGB特征，用于后续的异常分割任务。

EasyNet的算法流程：

输入图像对：
- 输入包括一对RGB图像和相应的深度图像。
异常图像生成：
- 异常生成器向深度图像添加噪声，生成异常图像。
- 同时生成相应的掩码，用于后续的重建任务。
特征提取与重建：
- MRN接收异常图像和原始图像，执行重建任务。
- 在重建过程中，MRN从不同层次提取多模态特征。
特征融合：
- MSN利用基于注意力的信息熵融合模块来融合RGB和深度特征。
- 该模块通过计算信息熵来决定是否融合深度特征。
异常分割：
- MSN输出异常分割图，标识出异常区域。
损失计算：
- 使用SSIM损失和MSE损失计算RGB图像的重建损失。
- 使用MSE损失计算深度图像的重建损失。
- 使用Focal损失计算像素分类任务的损失。
训练与优化：
- 通过优化上述损失函数来训练网络。
- 使用Adam优化器进行参数更新。
推理：
- 在推理阶段，EasyNet使用训练好的模型对新的图像进行异常检测。

pseudo-code

import torch
from torch import nn, optim

# 假设以下网络结构已经定义
class AnomalyGenerator(nn.Module):
    def forward(self, rgb_image, depth_image):
        # 添加噪声并返回异常图像
        pass

class MultiModalityReconstructionNetwork(nn.Module):
    def forward(self, aug_rgb_image, aug_depth_image):
        # 重建图像并提取多尺度特征
        pass

class MultiModalitySegmentationNetwork(nn.Module):
    def forward(self, multi_scale_features):
        # 融合特征并输出异常分割图
        pass

# 损失函数定义
def loss_function(anomaly_map, real_labels):
    # 计算并返回损失
    pass

# 训练模型
def train_model(model_mrn, model_msn, anomaly_generator, train_loader, epochs, optimizer_mrn, optimizer_msn, criterion):
    for epoch in range(epochs):  # 遍历每个训练周期
        for rgb_images, depth_images in train_loader:  # 遍历训练数据集
            # 生成异常图像
            aug_rgb_images, aug_depth_images = anomaly_generator(rgb_images, depth_images)
            
            # 重建图像并提取特征
            rec_rgb_images, rec_depth_images, multi_scale_features = model_mrn(aug_rgb_images, aug_depth_images)
            
            # 生成异常分割图
            anomaly_map = model_msn(multi_scale_features)
            
            # 计算损失
            loss = criterion(anomaly_map, real_labels)
            
            # 反向传播和优化
            optimizer_mrn.zero_grad()
            optimizer_msn.zero_grad()
            loss.backward()
            optimizer_mrn.step()
            optimizer_msn.step()

# 测试模型
def test_model(model_mrn, model_msn, anomaly_generator, test_loader):
    for rgb_images, depth_images in test_loader:  # 遍历测试数据集
        # 生成异常图像
        aug_rgb_images, aug_depth_images = anomaly_generator(rgb_images, depth_images)
        
        # 重建图像并提取特征
        rec_rgb_images, rec_depth_images, multi_scale_features = model_mrn(aug_rgb_images, aug_depth_images)
        
        # 生成异常分割图
        anomaly_map = model_msn(multi_scale_features)
        
        # 将异常图转化为最终结果
        final_results = post_process(anomaly_map)
        
    return final_results

# 后处理
def post_process(anomaly_map, threshold=0.5):
    binary_results = (anomaly_map > threshold).float()
    return binary_results

# 主程序
if __name__ == "__main__":
    # 初始化网络和优化器
    mrn = MultiModalityReconstructionNetwork()
    msn = MultiModalitySegmentationNetwork()
    anomaly_generator = AnomalyGenerator()
    
    optimizer_mrn = optim.Adam(mrn.parameters(), lr=0.002)
    optimizer_msn = optim.Adam(msn.parameters(), lr=0.002)
    criterion = nn.MSELoss()  # 以MSE损失为例

    # 加载数据
    train_loader = ...  # 训练数据加载器
    test_loader = ...   # 测试数据加载器

    # 训练模型
    train_model(mrn, msn, anomaly_generator, train_loader, epochs=50, optimizer_mrn=optimizer_mrn, optimizer_msn=optimizer_msn, criterion=criterion)

    # 测试模型
    final_results = test_model(mrn, msn, anomaly_generator, test_loader)