U-Net实现缺陷检测

先挖个坑,把之前做过的缺陷检测记录一下,以后有时间再来填坑

U-Net

U-Net网络发表于2015年,最开始是用于医学细胞图像分割,但是针对其他的分割问题,U-Net似乎也表现出了不错的性能[2]。该网络结构如下图所示,其整体的流程实际上是一个编码和解码(encoder-decoder)的过程。
U-Net实现缺陷检测_第1张图片
U-Net网络是一个经典的全卷积网络,输入为572×572大小的图片,论文中将网络左侧称为contracting path,是由卷积和Max Pooling构成的降采样操作。

contracting path主要由4个block组成,每一个block包含了三层卷积和一层最大池化操作,并且在每次降采样操作之后都会将Feature Map的个数乘二,最终尺寸大小为32×32的Feature Map。

网络右侧部分则被称为expansive path,跟contracting path一样,由4个block组成,每个block开始之前通过反卷积操作将特征图尺寸放大两倍,同时将特征通道数减半,然后和左侧对称的contracting path的特征图进行合并,由于左侧压缩路径和右侧扩展路径的Feature Map的尺寸不一样,U-Net通过将contracting path的Feature Map裁剪到和扩展路径相同尺寸的Feature Map来进行归一化。

最终得到的特征图尺寸大小为388×388,由于论文内为二分类任务,因此输出两张特征图。总结来说,U-Net的卷积层数量大约在20个左右,进行4次下采样,4次上采样。

算法实现

本次算法实现过程环境为tensorflow2.0,调用tf中已经继承好的keras来实现训练过程,未使用GPU加速。项目实际实现过程在Anaconda中搭建虚拟环境,然后在jupyter notebook下完成模型训练和验证过程。
U-Net实现缺陷检测_第2张图片

  1. 读取图片及预处理
    上图所示为本次使用的算法模型结构。图片和标签图片均存放在dataset文件夹下,开始时首先将图像和掩膜图进行读取,各自为1000张,在代码内部分别保存于images和annotaions中。由于本次需要对图片进行二值分割,因此训练过程中的掩膜图最最好为单通道二值图像。但是实际观察到训练使用的掩膜图片除0和255值之外还有一些离散像素的灰度值居于两者之间,因此需要对掩膜图先进行转码然后将其二值化为0和1分布的图,储存于anno变量中。完成后将images和anno进行随机打乱,但是要保证打乱之后的顺序依旧是依次对应的。之后可以开始制作dataset数据集。

  2. Dataset制作
    Dataset API是TensorFlow 1.3版本中引入的一个新的模块,主要服务于数据读取,构建输入数据的pipeline。此前,在TensorFlow中读取数据一般有两种方法:使用placeholder读内存中的数据和使用queue读硬盘中的数据。而Dataset API同时支持从内存和硬盘的读取,其具有大量处理数据的实用方法,且语法更加简洁易懂。Google官方给出的Dataset API中的类图如下所示,Dataset可以看作是相同类型元素的有序列表。在实际使用时,单个元素可以是向量,也可以是字符串、图片,元组tuple或者字典dict。
    U-Net实现缺陷检测_第3张图片
    首先使用tf.data.Dataset.from_tensor_slices((image, anno))加载数据,此时dataset里面的一个元素为(image, anno)。jupyter notebook命令行里面输入dataset查看内容,显示为

    然后调用dataset.map(load_dataset),map接收一个函数对象,Dataset中的每个元素都会被当作这个函数的输入,并将函数返回值作为新的Dataset。此处调用的load_dataset函数完成图片尺寸调整和归一化操作。调用完成后dataset里面储存格式如下:。内部对应1000条数据。

    之后将dataset数据集进行划分,随机取出800条数据作为训练集dataset_train,200条数据作为验证集dataset_val。对dataset_train内部数据进行transformation,使用tf中dataset的接口函数,batch()划分内部元素组成大小为32的batch(最后一组可能会小于32),shuttle()再次打乱内部元素,repeat()将整个序列重复多次。

  3. 训练模型并验证
    由于使用的是keras,因此首先建立模型,调用U-Net模型建立函数create_model()生成框架并且赋给model变量。然后model.compile()指定优化器为Adam,由于是二值分割,loss定义为binary_crossentropy,再在metrics中指定模型训练时的性能指标为accuracy和mIOU(mIOU通过调用tf.keras中自带的模块生成类进行测试)。

    最后调用model.fit()来进行U-Net模型的训练,由于是在CPU下运行,因此时间较长,整个训练过程持续大约18个小时。

  4. 训练结果
    获得loss和val_loss输出图像如下:
    U-Net实现缺陷检测_第4张图片
    在训练完成之后保存模型为unet_model.h5,从测试集中随机取出部分图像查看分割效果如下(注意将预测结果图的像素点从0~1还原为0, 255二值分布)。从左到右依次为:原图,给定的掩膜图和预测分割图。
    U-Net实现缺陷检测_第5张图片
    参考文献
    [1] Evan Shelhamer, Jonathan Long, Trevor Darrell. Fully Convolutional Networks for Semantic Segmentation[M]. IEEE Computer Society, 2017.
    [2] Ronneberger O , Fischer P , Brox T . U-Net: Convolutional Networks for Biomedical Image Segmentation[J]. 2015.
    [3] Simonyan K , Zisserman A . Very Deep Convolutional Networks for Large-Scale Image Recognition[J]. Computer Science, 2014.
    [4] He K , Zhang X , Ren S , et al. Deep Residual Learning for Image Recognition[C]// IEEE Conference on Computer Vision & Pattern Recognition. IEEE Computer Society, 2016.
    [5] 韩慧. 基于深度学习的工业缺陷检测方法研究[D]. 2019.
    [6] 刘聪. 基于卷积神经网络的微小零件表面缺陷检测技术研究[D]. 2019.

我一定会来填坑的,确信。。。

你可能感兴趣的:(机器学习,深度学习,tensorflow,python,神经网络,机器学习)