PR U-Net：工业缺陷分割中的特征融合和渐进监督策略

PR U-Net: Feature Fusion and Progressive Supervision Strategy in Industrial Defect Segmentation

摘要: 在工业缺陷检测任务中，由于缺陷与周围环境的相似性，准确识别图像中的缺陷区域是一项重要的挑战。本文提出了一种名为渐进细化(PR) U-Net的新型语义分割模型，该模型基于Swin U-Net架构，通过集成PR解码块和焦点模块(FM)来进行改进。PR解码块使模型能够逐步细化特征，并在每个解码阶段引导监督，而FM通过融合分割预测和特征融合机制，有助于更好地区分缺陷和非缺陷区域。大量实验结果表明，本文的PR U-Net在区分背景区域中的细微缺陷方面具有更高的准确性和鲁棒性，同时在推理时间方面也具有高效性，使其适用于各种场景下的工业缺陷检测任务。

Abstract: The paper presents a novel semantic segmentation model named Progressive Refining (PR) U-Net, designed to enhance industrial defect detection. The model builds upon the Swin U-Net architecture, incorporating PR Decoder Blocks and Focus Modules (FMs) to improve feature refinement and segmentation accuracy. The PR Decoder Blocks allow for progressive feature refinement, while FMs aid in distinguishing defects from non-defects by integrating segmentation predictions and feature fusion mechanisms. Extensive experiments demonstrate that PR U-Net outperforms existing models in accuracy, robustness, and efficiency in industrial defect detection tasks.

1. 引言

图像分割技术在工业领域中起着至关重要的作用，它有助于将图像划分为语义上有意义的区域。其应用使得工业企业能够实现自动化、提高产品质量、降低成本，并增强竞争力。

图像分割技术广泛应用于工业环境中[1]-[4]。例如，它广泛应用于工业机器人视觉系统中，提高了机器人准确识别和定位目标的能力[5]。此外，图像分割技术对工业自动化和监控也有重大贡献[6]。通过分析工厂内部环境和设备状态的分割图像，企业能够实现生产过程的实时监控和优化，从而及时发现和解决潜在问题，提高生产效率和资源利用率。值得注意的是，其最有价值的应用之一在于装配线生产中的缺陷检测[3] [7]。利用图像分割技术，工业企业能够快速、准确地识别产品表面的缺陷、瑕疵或其他不合格之处，从而实现产品质量的有效监控和控制。这对于提高产品质量、降低售后服务成本以及维护品牌声誉至关重要。

然而，工业领域中的图像分割技术面临着诸多挑战和问题。Liu [8]等人指出，矿物分割模型在部署、分割精度和区域重叠方面面临的挑战。同样，Hesamian [9]等人指出，深度学习分割模型在工业领域中面临的数据标注不足、类别不平衡以及过拟合导致的分割精度下降等问题。总体而言，分割结果存在大量的假阳性和假阴性，以及普遍的边界错误。

为应对这些挑战，本文提出了渐进式精细化U-Net (PR U-Net)。在Swin U-Net [10]架构的基础上，本文重新设计了解码器模块，加入了三个渐进式精细化解码块(PR解码块)。通过传播和监督分层损失，并逐步精细化分割结果，本文旨在提高分割结果的质量，同时减少边界错误。同时，在每个PR解码块中，本文引入了聚焦模块结构，以帮助模型更有效地区分缺陷区域和正常区域，从而减少假阳性和假阴性的发生。

2. 相关工作

2.1. 基于CNN的方法

最初，学者们通过将卷积神经网络(CNN)中的最后一层全连接层替换为卷积层，使模型输出特征图作为分割结果。考虑到不同层次之间尺寸和语义信息的差异，FCN [11]引入了上采样和跳跃连接，使高层特征能够直接添加到低层特征中，初步形成了特征金字塔的雏形。然而，FCN没有考虑全局特征，忽略了潜在有用的场景级语义上下文[12]。

2.2. 基于编码器–解码器的方法

在图像分割领域，编码器–解码器模型也是一类流行的深度学习模型。通用分割的编码器–解码器模型包括SegNet [13]、HRNet [14]等。其中，Noh [15]等人提出了一种基于反卷积的语义分割方法，使用VGG 16层网络的卷积层作为编码器，并通过反卷积网络生成像素级类别概率图。SegNet利用编码器–解码器架构，通过在解码器中使用池化索引进行非线性上采样实现图像分割。在特定应用领域，如医学图像处理[16]-[18]和工业图像[19]-[21]处理，U-Net [22]和V-Net [23]受到了广泛关注。U-Net由Ron博士提出，具有一个收缩路径和对称扩展路径，能够有效地从少量标注图像中学习。U-Net的许多变体在医学和工业领域发挥了重要作用。V-Net通过引入基于Dice系数的新目标函数，解决了前景和背景体素数量不平衡的问题。

2.3. 基于注意力机制的方法

注意力机制在计算机视觉领域得到了持续探索，使模型在处理过程中能够更专注于输入数据的相关部分。因此，在语义分割中应用注意力机制是一种自然趋势。Chen [24]等人提出了一种注意力机制，学习每个像素位置上多尺度特征的软权重。Huang [25]等人提出了一种相反于传统方法的方法，他们设计了一个反向注意力网络(RAN)，捕捉与目标类别无关的特征。为了增强特征，Fu [26]等人引入了双重注意力网络，利用自注意机制捕捉图像中的丰富上下文依赖关系。通过在扩张卷积网络中附加两种类型的注意力模块，该模型能够模拟空间和通道维度的语义相互依赖关系，以理解图像的语义结构并实现更好的分割结果。

除了这些方法外，还有一些工作探索了注意力机制在语义分割中的应用[27]-[29]。SwinU-Net结合了Swin Transformer和U-Net，是一个重要的里程碑。一方面，Swin Transformer [30]是一种强大的注意力机制网络，能够有效处理不同尺度的特征。另一方面，U-Net结构是一种经典的编码器–解码器结构，用于保留和融合不同分辨率的特征。Swin U-Net结合了两者的优势，在语义分割任务中表现出色，通过跨区域信息交换和多尺度特征融合实现更准确和稳定的图像分割。

2.4. 计算机视觉在工业中的应用

视觉技术在现代工业中发挥着至关重要的作用。其基本原理是使用摄像机捕捉目标物体的图像，然后使用计算机算法对这些图像进行处理和分析，实现目标物体的检测、识别和定位等功能。在制造业中，视觉技术广泛应用于产品工艺监控、缺陷检测等阶段。

以图像分割为例，它不仅是计算机视觉领域的重要研究方向，而且在工业领域的各种任务场景中得到了广泛应用。Song [31]等人提出了一种结合空间注意力和前馈神经网络的混合残差模块(SAFM)，替代编码器中的下采样层以提取缺陷特征，代表了工业缺陷检测领域的前沿方法。

总的来说，视觉技术在工业中的应用前景广阔且意义重大。随着计算机视觉和人工智能技术的发展，视觉技术为工业生产的智能化和自动化提供了强有力的支持，推动了工业领域的进步与发展。

3. 本文模型

在工业缺陷检测任务中，模型需要准确识别图像中的缺陷区域，但由于缺陷(正样本)与周围环境(负样本)的相似性，这一任务具有一定的难度。Swin U-Net作为一种结合了Swin Transformer和U-Net的语义分割模型，在通用语义分割任务中表现出色。然而，其特征提取和信息融合能力不足以有效地区分背景中的细小缺陷区域，从而降低了在工业缺陷检测任务中的效率。

为了解决这些问题，本文提出了渐进细化U-Net (PR U-Net)。见图1所示，PR U-Net保留了Swin U-Net的编码器部分，以保留其特征提取和侧向连接能力。为了进一步增强模型区分前景和背景的能力，本文重新设计了整个解码器部分，包含三个渐进细化解码块(PR解码块)。每个PR解码块包含两个焦点模块(FM)，用于在当前分辨率下细化特征图，并更好地区分目标和非目标区域。通过传递编码器中获得的特征并经过三个PR解码块，模型逐步提高其语义分割精度。此设计使PR U-Net能够更有效地利用特征细化和信息融合机制，从而提高其在工业缺陷检测任务中的性能。

Figure 1. PR-Unet structure

图1. PR-Unet结构

3.1. 焦点模块

为了更好地区分缺陷和非缺陷区域，减少假阳性和假阴性的数量，本文设计了一个焦点模块。见图2所示，焦点模块包含三个输入：深层语义特征$F_H$ 、浅层语义特征$F_L$ 与和分割预测$F_{seg}$ 焦点模块的目标是通过特征融合，提高最终分割结果的准确性。

步骤1：对深层语义特征$F_H$ 、和分割预测$F_{seg}$ 进行两倍上采样，得到上采样后的特征${F^{\prime }}_H$ 和${F^{\prime }}_{seg}$ 此步骤的目的是将低分辨率特征图上采样到与高分辨率特征图一致，以确保后续操作的一致性和有效性。

步骤2：在特征融合过程中，首先通过交叉熵操作，将浅层语义特征$F_L$ 与分割预测$F_{seg}$ 逐元素相乘，强调分割预测中的缺陷区域。然后，获得的特征与权重参数$\alpha$ 相乘，并从深层语义特征${F^{\prime }}_H$ 中减去，得到特征$F_1$ 。此步骤的目的是强调模型应关注的缺陷区域，并减少模型对背景区域的注意，从而减少假阳性。

如公式(1)所示：

$F_1=\alpha \times \text{CE Block}\left({F^{\prime }}_{seg}\odot F_L\right)-{F^{\prime }}_H$ (1)

Figure 2. Focus module

图2. 焦点模块结构

在第二个分支中，将浅层语义特征$F_L$ 与反向分割预测$1-{F^{\prime }}_{seg}$ 逐元素相乘，以排除分割预测中的缺陷区域。然后，同样通过交叉熵操作并乘以权重参数$\beta$ ，特征$F_2$ 。此步骤的目的是强调模型应忽略的区域，并减少模型对非缺陷区域的注意，从而减少假阴性。

如公式(2)所示：

$F_2=\beta \times \text{CE Block}\left(\left(1-{F^{\prime }}_{seg}\right)\odot F_L\right)$ (2)

步骤3：从两条分支获得的特征进行逐元素异或操作，得到焦点模块的输出$out$ 。

如公式(3)所示：

$out=F_1\oplus F_2$ (3)

通过这种设计，焦点模块有效地融合了不同层次的特征信息，并关注模型需要关注的区域，从而提高了模型在缺陷检测任务中的性能。

3.2. 渐进细化解码块

渐进细化解码块设计用于语义分割任务，通过逐步细化特征并融合上层解码器的分割结果，生成最终分割输出。见图3所示，渐进细化解码块包含两个分支，每个分支集成了一个焦点模块和一系列卷积块操作。

Figure 3. Progressively refine the decoding block structure

图3. 渐进细化解码块结构

步骤1：对输入特征进行两倍上采样，以保持分辨率一致。然后，渐进细化解码块由两个分支组成，每个分支包含一个焦点模块和一系列卷积块操作。

步骤2：在第一个分支中，输入分割结果$F_{seg}$ 、解码器特征$e_D$ 和卷积块操作后的编码器特征$e_E$ 进入焦点模块$F{M}_1$ ，得到特征$z_1$ 。然后，将$z_1$ 与卷积块操作后的编码器特征${e^{\prime }}_E$ 拼接，生成融合特征表示。最后，通过两次卷积操作，进一步细化融合特征，得到特征表示${z^{\prime }}_1$ 。

如公式(4) (5) (6)所示：

$z_1=F{M}_1\left(F_{seg},\text{CBlock}\left(e_E\right),e_D\right)$ (4)

$z_1\leftarrow \left[z_1;\text{CBlock}\left({e^{\prime }}_E\right)\right]$ (5)

${z^{\prime }}_1\leftarrow \text{CBlock}\left(\text{CBlock}\left(z_1\right)\right)$ (6)

在该分支中，本文强调了分割结果信息的及时融合，突出模型感兴趣的区域。通过将分割结果与解码器和编码器特征联合处理，模型能够更好地学习目标区域的特征表示。该分支增强了分割结果中目标区域的特征学习，提升了模型对缺陷等目标的表示能力和检测能力。

在第二个分支中，输入分割结果$F_{seg}$ 、上采样的编码器嵌入${e^{\prime }}_E$ 和解码器嵌入${e^{\prime }}_D$ 进入焦点模块$F{M}_2$ ，得到特征$z_2$ ，然后，将第一个分支生成的细化特征与第二个分支生成的特征进行级联操作，最后，通过三次卷积操作，进一步提取和融合级联后的特征，生成最终的输出$out$ 。

如公式(7) (8) (9)所示：

$z_2=F{M}_2\left(F_{seg},\text{CBlock}\left({e^{\prime }}_E\right),\text{CBlock}\left({e^{\prime }}_D\right)\right)$ (7)

$z=\left[{z^{\prime }}_1;z_2\right]$ (8)

$out=\text{CBlock}\left(\text{CBlock}\left(\text{CBlock}\left(z\right)\right)\right)$ (9)

最后，为了逐层引导$F_{seg}$ 的生成过程，在PR解码器块的每一层计算每层的$F_{seg}$ 与实际分割结果$y_{gt}$ 之间的损失，具体公式(10)如下：

${ℒ}_{total}={\displaystyle \sum _{l=1}^Lℒ}\left(F_{seg,l},y_{gt}\right)$ (10)

这里，L表示PR解码块的层数，$F_{seg,l}$ 表示第l层的分割结果，$y_{gt}$ 是实际分割结果。这个损失计算机制确保每层的分割结果受到监督信号的引导，有助于提高模型在语义分割任务中的学习效果，从而使模型能够更好地理解图像内容并生成准确的分割结果。

4. 实验

4.1. 数据集

所有实验均在MVTec AD异常检测(MVTec AD)数据集上进行，该数据集是检测工业制造中表面缺陷的广泛使用的数据集。MVTec AD数据集由MVTec Software GmbH创建，包含各种类型的高分辨率图像，这些图像包含裂纹、凹痕、划痕和斑点缺陷等各种表面异常。这些图像覆盖了包括塑料、纸张、纺织品和金属在内的多种材料。

MVTec AD数据集旨在促进工业表面缺陷检测算法的研究，作为评估算法在实际工业场景中性能的基准。每类图像分为训练集和测试集，训练集包含正常样本和异常样本，测试集仅包含异常样本。

该数据集的一个关键特性是其丰富的变异性，包括光照条件、视角、缺陷类型和表面材料的变化。这使得MVTec AD数据集成为评估算法鲁棒性和泛化能力的理想选择。

4.2. 参数设置

所有实验均在一个综合的软件和硬件环境中进行。操作系统为Ubuntu 20.04，主要库包括Python 3.9、CUDA 11.8和PyTorch 2.0。硬件配置为Intel i9 12900KF CPU和NVIDIA RTX 3090TI GPU，系统配备了32GB 3200MHz的内存。本文的PR U-Net模型训练了100个epoch，批处理大小为32。优化器为Adam，学习率为0.001。这些配置在整个实验过程中保持一致，确保了结果的可靠性和可重复性。

4.3. 消融实验

为了评估本文提出的创新的有效性和贡献，本文以Swin U-Net作为基准模型进行了消融研究。

见表1所示，消融实验结果表明，本文提出的PR U-Net在所有指标上均优于基准模型Swin U-Net。具体来说，PR U-Net在Dice系数、召回率、精确度、F1分数和IoU等指标上都表现出更高的数值。例如，在Dice系数方面，PR U-Net在负类上的表现为99.8593，显著高于Swin U-Net的99.6877，而在正类上的表现为96.5790，明显高于Swin U-Net的92.4085。此外，PR U-Net在平均指标上也有显著提升，其表现为98.2192，而Swin U-Net为96.0481。

见图4所示，在正常和异常区域交界处的分割准确性方面，PR U-Net模型显著优于Swin U-Net模型，这意味着本文实现了精细分割的研究目标。总体而言，PR U-Net在工业缺陷检测任务中表现出优越的性能，进一步验证了其在特征提取和信息融合方面的有效性和优越性。

Figure 4. Example of comparison of segmentation results in ablation experiments

图4. 消融实验分割结果对比示例

4.4. 与经典模型的比较

本文使用其他经典的图像分割模型在MVTec AD数据集上进行了实验。为了确保比较的严格性和公平性，这些模型都在相同的硬件配置和相同的方式下进行了训练。

见表2所示，PR U-Net在工业数据集上显著优于ViT [23]和TransUnet [27]这两个强大的分割模型。这证明了逐步细化分割在工业异常分割任务中的关键作用。

4.5. 与最新模型比较

在本节中，本文将PR U-Net模型与当前工业缺陷分割领域的最新模型SAFM Res UNet [31]在准确性和推理时间方面进行比较。

见表3所示，PR U-Net在这两个特定数据集上的图像分割能力均优于SAFM Res UNet。这意味着PR U-Net更能准确地识别工业产品中的细微缺陷和异常。这样的性能优势表明，PR U-Net在实际工业应用中更为可靠，特别是在需要极高精度和准确度的制造和质量控制领域。它还展示了PR U-Net在不同工业缺陷场景中的优秀泛化能力。本文还评估了前面提到的几种常见分割方法的推理时间。同样，所有测试均在相同配置下使用MVTec AD数据集进行，每个模型在该数据集上独立运行10次，以确定平均推理时间。为了更好地呈现实验结果，本文对推理时间进行了归一化处理。

见表4所示，PR U-Net在推理速度方面略优于SAFM Res UNet，这使得PR U-Net适用于需要快速和精确处理的工业应用。

Table 1. Comparison of ablation experimental result

表1. 消融实验结果对比

	类别	Dice	召回率	精确度	F1分数	IoU
Swin U-Net (基准)	负类	99.6877	99.6805	99.6949	99.6877	99.3774
Swin U-Net (基准)	正类	92.4085	92.5712	92.2465	92.4086	85.8885
Swin U-Net (基准)	平均	96.0481	96.1258	95.9707	96.0481	92.6329
PR U-Net	负类	99.8593	99.8528	99.8658	99.8593	99.719
PR U-Net	正类	96.5790	96.7325	96.4260	96.5700	93.3843
PR U-Net	平均	98.2192	98.2926	98.1459	98.2192	96.5517

Table 2. Comparison results of different classical models

表2. 不同经典模型的比较结果

	ViT	TransUnet	Swin U-Net	PR U-Net
Dice	93.2581	95.3832	96.0481	98.2192
F1分数	93.2546	95.3217	96.0481	98.2129
IoU	90.4587	91.2321	92.6329	96.5517

Table 3. PR U-Net and SAFM Res Unet based on Dice comparison

表3. PR U-Net与SAFM Res Unet基于Dice系数泛化性比较

	MVTec AD	KolektorSDD
SAFM Res UNet	97.8	95.1
Swin U-Net	96.0	94.5
PR U-Net	98.2	95.3

Table 4. Comparison of normalized inference time between different models

表4. 不同模型推理时间经归一化后的比较

	Vit	TransUnet	SAFM Res UNet	PR U-Net
平均推理时间	1	0.6	0.09	0

5. 结论

在本文中，本文提出了一种新型的语义分割模型——渐进细化(PR) U-Net，以提高工业缺陷检测的性能。PR U-Net基于Swin U-Net架构，通过集成PR解码块和焦点模块(FM)来进行改进。PR解码块允许逐步细化特征，并在每个解码阶段进行监督，而焦点模块通过特征融合和分割预测，增强了模型区分缺陷本文和非缺陷区域的能力。

本文在MVTec AD和KolektorSDD两个广泛使用的工业缺陷检测数据集上进行了实验，结果表明，PR U-Net在所有评估指标上均显著优于现有的基准模型Swin U-Net以及其他经典的分割模型，如ViT和TransUnet。此外，PR U-Net在准确性、召回率、精确度、F1分数和IoU等方面均表现出色，特别是在细小缺陷区域的分割能力上，取得了显著的提升。

通过消融研究，本文进一步验证了PR解码块和焦点模块在提高模型性能方面的有效性。这些模块的设计不仅提高了特征提取和信息融合的能力，还有效地减少了假阳性和假阴性的发生，从而提高了分割的精度和鲁棒性。

本文还对比了PR U-Net与最新的工业缺陷检测模型SAFM Res UNet，结果显示PR U-Net在多个数据集上的性能均优于SAFM Res UNet，证明了本文模型在工业应用中的有效性和广泛适用性。此外，PR U-Net在推理时间上也表现出色，适合需要快速处理的工业场景。

总的来说，PR U-Net在工业缺陷检测任务中表现出优越的性能和高效性，证明了其在实际应用中的潜力。未来的研究可以进一步优化模型结构，并探索其在其他图像分割任务中的应用，进一步提升模型的通用性和适应性。

本文计划未来的工作将集中在以下几个方面：首先，进一步优化PR U-Net的结构，提高其在更大规模数据集上的训练和推理效率；其次，探索更多种类的焦点模块，以进一步增强模型在复杂背景下的分割性能；最后，将PR U-Net应用于其他领域的图像分割任务，如医学图像分割和卫星图像分割，以验证其通用性和适用性。

综上所述，PR U-Net是一种有效且高效的工业缺陷检测模型，其创新性设计在多个实验中表现出色，具有广泛的应用前景和研究价值。