U-Net实现缺陷检测

先挖个坑,把之前做过的缺陷检测记录一下,以后有时间再来填坑
U-Net
U-Net网络发表于2015年,最开始是用于医学细胞图像分割,但是针对其他的分割问题,U-Net似乎也表现出了不错的性能[2] 。该网络结构如下图所示,其整体的流程实际上是一个编码和解码(-)的过程 。
U-Net网络是一个经典的全卷积网络,输入为572×572大小的图片,论文中将网络左侧称为 path,是由卷积和Max 构成的降采样操作 。
path主要由4个block组成,每一个block包含了三层卷积和一层最大池化操作,并且在每次降采样操作之后都会将 Map的个数乘二,最终尺寸大小为32×32的 Map 。
网络右侧部分则被称为 path,跟 path一样,由4个block组成,每个block开始之前通过反卷积操作将特征图尺寸放大两倍,同时将特征通道数减半,然后和左侧对称的 path的特征图进行合并,由于左侧压缩路径和右侧扩展路径的 Map的尺寸不一样,U-Net通过将 path的 Map裁剪到和扩展路径相同尺寸的 Map来进行归一化 。
最终得到的特征图尺寸大小为388×388,由于论文内为二分类任务,因此输出两张特征图 。总结来说,U-Net的卷积层数量大约在20个左右,进行4次下采样,4次上采样 。
【U-Net实现缺陷检测】算法实现
本次算法实现过程环境为.0,调用tf中已经继承好的keras来实现训练过程,未使用GPU加速 。项目实际实现过程在中搭建虚拟环境,然后在 下完成模型训练和验证过程 。
读取图片及预处理

U-Net实现缺陷检测

文章插图
上图所示为本次使用的算法模型结构 。图片和标签图片均存放在文件夹下,开始时首先将图像和掩膜图进行读取,各自为1000张,在代码内部分别保存于和中 。由于本次需要对图片进行二值分割,因此训练过程中的掩膜图最最好为单通道二值图像 。但是实际观察到训练使用的掩膜图片除0和255值之外还有一些离散像素的灰度值居于两者之间,因此需要对掩膜图先进行转码然后将其二值化为0和1分布的图,储存于anno变量中 。完成后将和anno进行随机打乱,但是要保证打乱之后的顺序依旧是依次对应的 。之后可以开始制作数据集 。
制作
API是 1.3版本中引入的一个新的模块,主要服务于数据读取,构建输入数据的 。此前,在中读取数据一般有两种方法:使用读内存中的数据和使用queue读硬盘中的数据 。而 API同时支持从内存和硬盘的读取,其具有大量处理数据的实用方法,且语法更加简洁易懂 。官方给出的 API中的类图如下所示,可以看作是相同类型元素的有序列表 。在实际使用时,单个元素可以是向量,也可以是字符串、图片,元组tuple或者字典dict 。
首先使用tf.data..((image, anno))加载数据,此时里面的一个元素为(image, anno) 。命令行里面输入查看内容,显示为 。
然后调用.map(),map接收一个函数对象,中的每个元素都会被当作这个函数的输入,并将函数返回值作为新的 。此处调用的函数完成图片尺寸调整和归一化操作 。调用完成后里面储存格式如下:
。内部对应1000条数据 。
之后将数据集进行划分,随机取出800条数据作为训练集,200条数据作为验证集 。对内部数据进行,使用tf中的接口函数,batch()划分内部元素组成大小为32的batch(最后一组可能会小于32),()再次打乱内部元素,()将整个序列重复多次 。
训练模型并验证
由于使用的是keras,因此首先建立模型,调用U-Net模型建立函数()生成框架并且赋给model变量 。然后pile()指定优化器为Adam,由于是二值分割,loss定义为,再在中指定模型训练时的性能指标为和mIOU(mIOU通过调用tf.keras中自带的模块生成类进行测试) 。