深度学习与图像处理实战-基于 U-Net 的工业缺陷检测

基于U-Net的工业缺陷检测深度学习与图像处理实战知识要点10.1U-Net简介10.2数据集介绍及处理10.3主干网络搭建与训练10.4模型训练结果测试目录10.1U-Net简介FCN开创了深度学习在图像分割领域应用的先河，其变种U-Net作为图像分割网络的一种，在医学领域以及工业领域应用十分广泛。首先要理解一个概念：U-Net不是一个特定的网络，它只是一个网络结构的代号而已。其网络结构如图所示。U-Net结构10.1U-Net简介U-Net结构的形状呈U字，整体结构是先编码（下采样，左侧网络），再解码（上采样，右侧网络），接着回归到跟原始图像一样大小的像素的分类。U-Net结构的左半部分是特征提取部分，U-Net经过4次下采样进行特征提取，并且它和FCN一样除去了全连接层，在特征提取的最后一层直接进行上采样，并与前面相同尺度的特征进行融合，这就是U-Net中上采样的部分。10.1U-Net简介下采样是通过2×2的最大池化层来进行的，下采样之间是两个卷积层，这里的卷积使用无填充模式，所以在卷积过程中图像的大小是会减小的。这会造成一个问题，即在进行特征融合的时候图像大小不一致，所以在U-Net结构图中有一个copy&crop操作（剪裁操作），crop就是为了将大小进行裁剪的操作。但若在卷积的时候使用的是填充模式，就无须进行crop操作。10.2数据集介绍及处理本章数据集使用的是Kaggle上的数据，经解压后共有10个类别，如图所示，分别对应不同的花纹缺陷，如磨损、白点、多线等。数据类别数据详情并且，每个类别里面有Train、Test两个文件夹，分别存放着图片数据和标签数据，但并不是每张图片都有对应的标签，数据详情如图所示。10.2数据集介绍及处理这样的数据处理起来比较复杂，本章中使用的方式是将它们的路径一

一读取出来，然后进行随机打乱，按比例划分，最后把它们写入TXT文件中，后续只需要读取这个文件中的路径即可。具体步骤如下。首先，新建一个writeline.py文件，写入如下代码。01OPTION数据集处理importosimportglobimportrandom#读取有标签数据的图片defwrite_lines(split=0.9):lines=[]foriinrange(1,11):10.2数据集介绍及处理label_dir=r'\DAGM_KaggleUpload\Class%s\Train\Label'%i#label_paths=[os.path.join(label_dir,p)forpinos.listdir(label_dir)]#读取扩展名为PNG的图片label_paths=glob.glob(label_dir+'\*.PNG')#print(path)forlabel_pathinlabel_paths:name=label_path.split('\\')[-1].split('_')[0]img_path='\\'.join(label_dir.split('\\')[:-1])img_path=img_path+'\%s.PNG'%namelines.append('%s%s\n'%(img_path,label_path))#随机打乱random.shuffle(lines)train_len=int(len(lines)*split)train_lines=lines[:train_len]test_lines=lines[train_len:]#写入文件withopen('train.txt','w',encoding='utf8')asf:f.writelines(train_lines)withopen('test.txt','w',encoding='utf8')asf:f.writelines(test_lines)if__name__=='__main__':write_lines()10.2数据集介绍及处理可以发现在当前目录中多出了两个TXT文件，分别是test.txt、train.txt。里面存放的数据结构为“数据图片路径+空格+标签图片路径”，如图所示。02OPTION运行程序存放的数据10.2数据集介绍及处理读取数据的时候只需要读取这两个TXT文件中的数据即可。新建一个load_data.py文件，定义一个readline函数用来读取路径，代码如下。03OPTION读取数据importnumpyasnpfromPILimportImage#读取TXT文件中的数据defreadline(path):withopen(path,'r',encoding='utf8')asf:lines=f.readlines()img_lists=[]label_list=[]forlineinlines:line=line.strip('\n')data,label=line.split('')img_lists.append(data)label_list.append(label)returnimg_lists,label_list10.2数据集介绍及处理接着，定义一个数据处理函数get_data用来加载图片，并对标签进行处理，代码如下。04OPTION数据处理、加载图片、缩放图片#数据处理defget_data(x,y,size):#加载图片data=cv2.imread(x)label=Image.open(y)h,w=size,size#缩放图片label=label.resize((h,w))image_data=cv2.resize(data,(size,size))/255.

label=np.array(label)label[label>0]=1.returnimage_data,label10.2数据集介绍及处理代码如下。05OPTION组装一个生成器函数gen_datadefgen_data(x_data,y_data,batch_size,size=224):whileTrue:data=[]label=[]forindex,image_pathinenumerate(x_data):data_,label_=get_data(image_path,y_data[index],size)data.append(data_)label.append(label_)iflen(data)==batch_size:data=np.array(data).reshape(-1,size,size,3)label=np.array(label).reshape(-1,size,size,1)#print(label.shape,data.shape)yielddata,labeldata=[]label=[]10.2数据集介绍及处理最后，为了确认数据与标签是否一

一匹配，还可以在main函数中写入如下代码。06OPTION使用main函数if__name__=='__main__':color=np.array([(0,0,0),(0,255,0)])x_train,y_train=readline('train.txt')gan=gen_data(x_train,y_train,1)foriinrange(100):data,label=next(gan)data=np.array(data*255.,dtype='uint8').reshape(224,224,3)image=Image.fromarray(data)

label=label.reshape(224,224).astype('int32')color_image=np.array(corlor)[label.ravel()]\.reshape(224,224,3).astype('uint8')label=Image.fromarray(color_image)img=cv2.addWeighted(data,1,color_image,0.5,0)cv2.imshow('s',img)cv2.waitKey(0)10.2数据集介绍及处理运行结果如图所示。可以发现，这些数据的标注都比较粗糙，并没有过多地追求细节，而是用一个椭圆框住了缺陷内容，但这不影响后续的训练。数据与标签匹配10.3主干网络搭建与训练接着，新建一个model.py文件，写入如下代码，建立一个U-Net模型结构。01OPTION建立一个U-Net模型结构importtensorflow.kerasaskfromtensorflow.keras.layersimport*defu_net(inpt):#基础卷积结构：Conv+BN+ReLUdefconv(x,filters,kernel=3,strides=1,padding='same'):x=Conv2D(filters,kernel_size=kernel,strides=strides,padding=padding)(x)x=BatchNormalization()(x)x=Activation('relu')(x)returnx10.3主干网络搭建与训练inpt=k.Input(inpt)x=conv(inpt,64)x=conv(x,64)x1=x#第一层特征向量x=MaxPool2D()(x)#128x=conv(x,128)x=conv(x,128)x2=x#第二层特征向量x=MaxPool2D()(x)#64x=conv(x,256)x=conv(x,256)x3=x#第三层特征向量x=MaxPool2D()(x)#32x=conv(x,512)x=conv(x,512)x4=x#第四层特征向量x=MaxPool2D()(x)#16x=conv(x,1024)x=conv(x,1024)#上采样，并与第四层特征向量融合x=UpSampling2D()(x)#32x=concatenate([x,x4])x=conv(x,512)x=conv(x,512)#与第三层特征向量融合x=UpSampling2D()(x)#64x=concatenate([x,x3])x=conv(x,256)x=conv(x,256)#与第二层特征向量融合x=UpSampling2D()(x)#128x=concatenate([x,x2])10.3主干网络搭建与训练x=conv(x,128)x=conv(x,128)#与第一层特征向量融合x=UpSampling2D()(x)x=concatenate([x,x1])x=conv(x,64)x=conv(x,64)#通道缩减成类别个数，这里是二分类，并且使用Sigmoid进行分类x=Conv2D(1,kernel_size=(1,1),strides=1,padding='same')(x)x=Activation('sigmoid')(x)model=k.models.Model(inpt,x)model.summary()returnmodel10.4模型训练结果测试训练完成之后，可以使用如下代码进行测试。importtensorflowastfimportnumpyasnpfromFCN.Mobile_UNet.load_dataimportreadline,gen_dataimportcv2#不同类别对应的颜色corlor=np.array([(0,0,0),(0,255,0)])#加载模型model=tf.keras.models.load_model(r'mb_unet.h5')x_tests,_=readline('test.txt')forimage_pathinx_tests:#数据处理img=cv2.imread(image_path)img=cv2.resize(img,(224,224))image=img.copy()img=img.reshape(-1,224,224,3)/255.p=model.predict(img)#将预测的数据进行分类p[p>0.5]=1p[p<0.5]=0

p=p.reshape(224,224).astype('uint8')color_image=np.array(corlor)[p.ravel()]\.reshape(224,224,3).astype('uint8')img=cv2.addWeighted(image,1,color_image,0.5,0)cv2.imshow('s',img)cv2.waitKey(0)10.4模型训练结果测试测试程序运行结果如图所示。测试程序运行结果效果较差的情况是因为图片数据仅有941张，而训练类别却有10个，平均下来每个类别的图片不到95张，数据相对来说是缺乏的。因此需要增加数据量10.4模型训练结果测试

数据增强的方式有很多种，在本书的第5章已经有比较详细的介绍。这里引入随机数，对图片进行随机缩放、翻转和颜色空间变换这3种数据增强。将get_data函数修改为get_random_data函数，并写入如下代码。#获得a～b的随机数defrand(a=0,b=1):returnnp.random.rand()*(b-a)+adefget_random_data(x,y,size,jitter=.3,hue=.1,sat=1.5,val=1.5):data=Image.open(x)label=Image.open(y)h,w=size,size#对图片进行随机缩放new_ar=w/h*rand(1-jitter,1+jitter)/rand(1-jitter,1+jitter)scale=rand(.25,2)ifnew_ar<1:nh=int(scale*h)nw=int(nh*new_ar)else:nw=int(scale*w)nh=int(nw/new_ar)data=data.resize((nw,nh),Image.BICUBIC)label=label.resize((nw,nh))10.4模型训练结果测试#随机坐标dx=int(rand(0,w-nw))dy=int(rand(0,h-nh))#创建一张新的图片，并把缩放后的图片放到新图片的随机坐标中new_data=Image.new('RGB',(w,h),(128,128,128))new_label=Image.new('RGB',(w,h),(0,0,0))new_data.paste(data,(dx,dy))new_label.paste(label,(dx,dy))data=new_datalabel=new_label

#图片翻转flip=rand()<.5ifflip:data.transpose(Image.FLIP_LEFT_RIGHT)label.transpose(Image.FLIP_LEFT_RIGHT)

#HSV颜色空间变换hue=rand(-hue,hue)sat=rand(1,sat)ifrand()<.5else1/rand(