U-Net实现缺陷检测

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

U-Net

U-Net网络发表于2015年，最开始是用于医学细胞图像分割，但是针对其他的分割问题，U-Net似乎也表现出了不错的性能[2]。该网络结构如下图所示，其整体的流程实际上是一个编码和解码（encoder-decoder）的过程。

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下完成模型训练和验证过程。

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。
   
   首先使用tf.data.Dataset.from_tensor_slices((image, anno))加载数据，此时dataset里面的一个元素为(image, anno)。jupyter notebook命令行里面输入dataset查看内容，显示为<TensorSliceDataset shapes: ((), ()), types: (tf.string, tf.string)>。

   然后调用dataset.map(load_dataset)，map接收一个函数对象，Dataset中的每个元素都会被当作这个函数的输入，并将函数返回值作为新的Dataset。此处调用的load_dataset函数完成图片尺寸调整和归一化操作。调用完成后dataset里面储存格式如下：<ParallelMapDataset shapes: ((256, 256, 3), (256, 256, 1)), types: (tf.float32, tf.int32)>。内部对应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输出图像如下：
   
   在训练完成之后保存模型为unet_model.h5，从测试集中随机取出部分图像查看分割效果如下（注意将预测结果图的像素点从0~1还原为0, 255二值分布）。从左到右依次为：原图，给定的掩膜图和预测分割图。
   
   参考文献
   [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.

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