深度学习实现工业零件的缺陷检测

介绍

工业零件在制造完成的时候，往往需要去检测其完整性和功能性。如下图所示，从左上到右下，分别是擦花、漏底、碰凹、凸粉。本篇博文主要讲解如何去识别这四类图像，所用框架为keras-2.1.6+tensorflow-1.7.0+GTX1060。

数据集

所获得的数据不是太多，一共250幅图像，其中擦花图像29幅、漏底图像140幅、碰凹图像20幅、凸粉图像61幅。可以发现，这样的数据对深度卷积网络来说是远远不够的，而且数据分布是极不对称的。而且对特征是否明显来说，擦花特征是最不明显的，在训练的时候更要进行关注。

图像增强

对每类图像进行增强，使其扩展到300幅左右。这里使用ImageDataGenerator来完成，主要包括图像的旋转、平移、错切等。

from keras.preprocessing.image import ImageDataGenerator
import numpy as np
import numpy as np
import keras
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.optimizers import SGD
import tensorflow as tf
from PIL import Image
from image_handle import *
import os
import cv2


batch_size=10
epoch_size=50
data = np.empty((250,800, 800,3), dtype="float32")
image_resize=[]
image_vector=[]
image_label=[]
image_label_list=['擦花','漏底','碰凹','凸粉']  #漏底占的比例太大了，导致后面的预测全部都偏向该物体
image_dir='F:/Data/guangdong_round1_train1_20180903/'
j=0
k=0

for fn in os.listdir(image_dir):
    image_name=get_image_name(fn)
    for i in range(0,4):
        if image_label_list[i]==image_name:
            image_name=i
    image = Image.open(image_dir+fn,'r')#原始图片过大
    image_down=image.resize((800, 800), Image.ANTIALIAS)
    arr = np.asarray(image_down, dtype="float32")
    data[j, :, :, :] = arr
    j=j+1
    image_label.append(image_name)

#准备好数据和标签送进VGG网络中
x_train=data[231:250]
y_train = keras.utils.to_categorical(image_label,4)[231:250]#转换成one-hot向量

print(len(x_train))
print(len(y_train))
print(image_label)

generator = ImageDataGenerator(
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.1,
        zoom_range=0.1,
        horizontal_flip=True,
    )

for batch in generator.flow(x_train,
                          batch_size=100,
                          save_to_dir='C:/Users/18301/Desktop/image_data',#生成后的图像保存路径
                          save_prefix='碰凹2018',
                          save_format='jpg'):
    k += 1
    if k > 10: #这个10指出要扩增多少个数据，最终的数据为x_train的大小与10的乘积，可根据要求自行更改
        break  # otherwise the generator would loop indefinitely

增强后的图像：

训练-Training

主程序，main.py

import numpy as np
impor