当前位置：首页 > article >正文

【2024工业图像异常检测文献】SuperSimpleNet: 统一无监督和监督式学习检测快速可靠的表面缺陷检测方法

article 2024/11/1 10:24:40

SuperSimpleNet: Unifying Unsupervised and Supervised Learning for Fast and Reliable Surface Defect Detection

1、Background

提出的SuperSimpleNet方法，是一种从SimpleNet演变而来的创新的判别模型。

【2023工业图像异常检测文献】SimpleNet: 简单的图像异常检测和定位网络

SuperSimpleNet在SimpleNet的基础上进行了一系列的改进和创新，以下是主要的改进点和创新之处：

特征上采样：SuperSimpleNet引入了特征上采样步骤，通过增加特征图的分辨率来提高对小缺陷的检测能力。这是对SimpleNet的一个直接改进，有助于提升模型对细节的捕捉能力。
分类头部：新增了一个分类头部，用于生成图像级别的异常分数。这个分数反映了整个图像中异常的置信度，有助于减少误报并提高对小异常的检测率。
合成异常生成：在无监督设置中，SuperSimpleNet使用Perlin噪声和高斯噪声生成合成异常，这些异常被限制在非异常区域，而在监督设置中，合成异常与真实异常结合使用，以增强模型对异常的泛化能力。
损失函数和训练策略：SuperSimpleNet采用了截断的L1损失和焦点损失，以优化模型的训练过程。此外，还引入了学习率调度器和梯度调整策略，以提高训练的稳定性和模型的性能。
无监督和监督学习的统一：SuperSimpleNet能够灵活地在无监督和监督设置中进行训练，这使得模型能够适应不同的训练条件，并充分利用所有可用的训练数据。

2、Method

SuperSimpleNet算法流程：

1. 特征提取（Feature Extraction）

预训练网络：使用预训练的卷积神经网络（例如WideResNet50），从输入图像中提取特征。
上采样：为了提高小缺陷的检测能力，对提取的低分辨率特征图进行上采样，使其具有更高的分辨率。
池化：通过局部平均池化来捕捉特征图邻近区域的信息，增强特征的上下文感知能力。

2. 特征适配（Feature Adaptation）

线性层：通过一个线性层（特征适配器）进一步调整特征，使其更适合后续的异常检测任务。

3. 合成异常生成（Synthetic Anomaly Generation）

Perlin噪声：使用Perlin噪声生成一个二值掩码，该掩码决定了在图像的哪些区域添加噪声。
高斯噪声：在无监督设置中，高斯噪声被添加到Perlin噪声掩码指定的所有区域。在监督设置中，高斯噪声被添加到非异常区域，而真实异常区域不添加噪声。

4. 分割头和分类头（Segmentation and Classification Heads）

分割头：负责预测输入图像的异常掩码（Mo），即识别图像中的异常区域。
分类头：使用分割头的输出和扰动后的特征图作为输入，产生一个异常分数（s），表示图像中异常的置信度。

5. 训练过程（Training Process）

损失函数：使用截断的L1损失和焦点损失来训练分割头，同时使用焦点损失来训练分类头。
监督信号：在监督设置中，模型的训练会受到真实异常掩码和异常标签的监督。在无监督设置中，模型仅使用合成异常进行训练。

6. 推理过程（Inference Process）

省略异常生成：在推理时，不生成合成异常，直接使用特征适配器的输出来预测异常掩码和异常分数。

在这里插入图片描述

pseudo-code

import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义SuperSimpleNet模型
class SuperSimpleNet(nn.Module):
    def __init__(self):
        super(SuperSimpleNet, self).__init__()
        # 特征提取器：预训练的卷积网络
        self.feature_extractor = PretrainedCNN()
        # 特征适配器：线性层
        self.feature_adaptor = nn.Linear(in_features, out_features)
        # 分割头：用于预测异常掩码
        self.segmentation_head = SegmentationHead()
        # 分类头：用于预测异常分数
        self.classification_head = ClassificationHead()

    def forward(self, x, is_training=False):
        # 特征提取
        features = self.feature_extractor(x)
        # 特征上采样和池化
        features = self.upsample_and_pool(features)
        # 特征适配
        adapted_features = self.feature_adaptor(features)
        
        if is_training:
            # 合成异常生成
            perturbed_features = self.generate_synthetic_anomalies(adapted_features)
        else:
            perturbed_features = adapted_features
        
        # 分割头预测异常掩码
        anomaly_mask = self.segmentation_head(perturbed_features)
        # 分类头预测异常分数
        anomaly_score = self.classification_head(perturbed_features)
        
        return anomaly_mask, anomaly_score

    def generate_synthetic_anomalies(self, features):
        # 使用Perlin噪声生成异常掩码
        perlin_mask = self.generate_perlin_noise_mask(features.size(0), features.size(2), features.size(3))
        # 添加高斯噪声生成合成异常
        perturbed_features = features + self.gaussian_noise(perlin_mask)
        return perturbed_features

    def upsample_and_pool(self, features):
        # 上采样特征图
        upsampled_features = F.interpolate(features, scale_factor=2, mode='bilinear')
        # 应用池化
        pooled_features = F.avg_pool2d(upsampled_features, kernel_size=3, stride=1, padding=1)
        return pooled_features

# 预训练CNN特征提取器
class PretrainedCNN(nn.Module):
    def __init__(self):
        super(PretrainedCNN, self).__init__()
        # 这里可以加载一个预训练的CNN模型，例如ResNet

    def forward(self, x):
        # 提取特征
        return features

# 分割头
class SegmentationHead(nn.Module):
    def __init__(self):
        super(SegmentationHead, self).__init__()
        # 定义分割头的层

    def forward(self, x):
        # 预测异常掩码
        return anomaly_mask

# 分类头
class ClassificationHead(nn.Module):
    def __init__(self):
        super(ClassificationHead, self).__init__()
        # 定义分类头的层

    def forward(self, x):
        # 预测异常分数
        return anomaly_score

# 生成Perlin噪声掩码
def generate_perlin_noise_mask(batch_size, height, width):
    # 生成噪声掩码
    return perlin_mask

# 生成高斯噪声
def gaussian_noise(mask):
    # 生成高斯噪声
    return noise