基于U-Net+残差网络的语义分割缺陷检测

一、介绍

基于深度学习的缺陷检测主要集中在场景识别、object detection等方法，近年来，推出了一系列优秀的语义分割模型，比如SegNet、FCN、U-Net等。语义分割模型被广泛的应用到场景识别、道路分割等领域中，缺陷检测领域也逐渐接受该模型。

本博文主要关注和实现基于U-Net的语义分割模型。

二、U-Net

如下图所示，这是U-Net的核心思想，从编码到解码，编码过程就是下采样，解码时，把同级的编码特征进行拼接，一起上采样。

详细算法请参考：https://blog.csdn.net/natsuka/article/details/78565229

U-Net由于把编码特征加入到解码特征中，可以在小数据量下取得很好效果，但是有个问题，U-Net网络不够深，对特征的表达不够准确，残差网络之所以取得这么好的分类效果其实就是网络足够深，一般来说网络更深，表达能力更强。

下面是改进的U-Net，其中ResNet Block数量为6。

三、实验数据

如下图所示，对缺陷的地方直接赋值像素为255，其它区域像素值保持不变，这样就可以生成标签数据。在代码中，对训练图像和标签图像分别进行0~1处理，对输入图像直接除以255，对标签图像，生成二进制图像，像素值为255，则置1，否则置0。

四、测试结果

五、算法改进

实验表明，如果只使用U-Net+残差模块，对大的缺陷区域分割效果比较好，但是对小的区域，效果则不好，为此，在原模型的基础上引入生成对抗网络来监督。这一网络结构利用了博弈对抗的思想，为其中生成式模型添加了有效的监督反馈，可以进一步提高缺陷分割网络对于图像缺陷区域标记的能力。

1. 对于作为输入的一组图片(包含原图和标记图)，首先训练判别式网络，此时生成式模型参数固定。训练的目的在于使判别式模型具有可以分辨输入的图像属于真实人工标记图还是生成网络生成的图像，于是训练过程如下:
   a、把输入的一对图片中的原始工件图x和人工标记图y同时输入到判别式网络D中（输入为两张图像），因为此时的y是x的真实对应，所以计算判别式网络的损失函数时标签应该为真，然后以反向传播算法更新判别式网络;
   b、把数据对中的工件原图x输入到生成式网络G中，得到缺陷分割生成图G (x)，然后x与G (x)同时输入到判别式网络D中，此时的G (x)为生成图像，不是x的真实对应图，所以计算判别式网络损失函数时标签应该设为假，同样地以反向传播算法更新判别式网络。
2. 然后训练生成式网络，同样地判别式模型维持参数不变，训练过程如下:
   a、把数据对中原始工件图x输入到生成式网络G，得到缺陷分割图G (x)，计算G (x)与人工标记图y的L1损失函数值(图像空间距离)，记为$L_G$;
   b、把G (x)与原始工件图像x输入到判别式网络D，由于训练生成式网络G的目的是要求其能生成的分割图接近真实标记图，所以此时计算判别式网络的损失函数应该设标签为真。经过训练后的判别式网络具备判断输入图像是真实标记图还是生成图像的能力，所以此时它的损失函数值能反映输入图像G (x)与真实图像的相似程度，记为$L_D$;
   c、更新生成式网络G的参数时，把生成图与真实标记图的空间距离$L_G$与判别式网络计算所得的相似程度$L_D$同时作为反向传播算法的代价函数，公式表示为$L_{total}=L_G+\lambda L_D$，其中兄为判别式网络损失函数权重，用于决定判别式网络对生成式网络的监督反馈强度，当λ=0时，等价于传统的语义分割网络训练方法。最后利用反向传播算法更新生成式网络G的参数。
3. 对于每一对训练样本，重复(1) (2)两个步骤，直到所有训练样本都参与到训练中，则完成了一次训练。当训练次数达到设定量时，训练结束。

算法的改进在于此，还没有做好，感兴趣的话，可以自己实现下~