Keras快速搭建简易缺陷检测网络

本文使用Keras来快速搭建一个简单的语义分割网络来进行单类别的瑕疵检测。

from keras import backend as K
from keras.models import Sequential
from keras.layers import Dense,Conv2D,MaxPool2D,Flatten,Activation,UpSampling2D
import numpy as np
import cv2
import os
import tensorflow as tf
import skimage.io as io
import skimage.transform as trans
from keras.preprocessing.image  import ImageDataGenerator
from keras.optimizers import Adam

input_shape=(1280,512)
def adjustData(img,mask,flag_multi_class,num_class):
    if(flag_multi_class):
        img = img / 255#img shape(2,512,256,1)
        mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]
        new_mask = np.zeros(mask.shape + (num_class,))#new_mask shape(2,512,256,num_class)
        for i in range(num_class):
            #for one pixel in the image, find the class in mask and convert it into one-hot vector
            #index = np.where(mask == i)
            #index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)
            #new_mask[index_mask] = 1
            new_mask[mask == i,i] = 1
        new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1]*new_mask.shape[2],new_mask.shape[3])) #if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))
        mask = new_mask
    elif(np.max(img) > 1):
        img = img / 255
        mask = mask /255
        mask[mask > 0.5] = 1
        mask[mask <= 0.5] = 0
    return (img,mask)



def trainGenerator(batch_size,train_path,image_folder,mask_folder,aug_dict,image_color_mode = "grayscale",
                    mask_color_mode = "grayscale",image_save_prefix  = "image",mask_save_prefix  = "mask",
                    flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = input_shape,seed = 1):
    '''
    can generate image and mask at the same time
    use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same
    if you want to visualize the results of generator, set save_to_dir = "your path"
    '''
    image_datagen = ImageDataGenerator(**aug_dict)
    mask_datagen = ImageDataGenerator(**aug_dict)
    image_generator = image_datagen.flow_from_directory(
        train_path,
        classes = [image_folder],
        class_mode = None,
        color_mode = image_color_mode,
        target_size = target_size,
        batch_size = batch_size,
        save_to_dir = save_to_dir,
        save_prefix  = image_save_prefix,
        seed = seed)
    mask_generator = mask_datagen.flow_from_directory(
        train_path,
        classes = [mask_folder],
        class_mode = None,
        color_mode = mask_color_mode,
        target_size = target_size,
        batch_size = batch_size,
        save_to_dir = save_to_dir,
        save_prefix  = mask_save_prefix,
        seed = seed)
    train_generator = zip(image_generator, mask_generator)
    for (img,mask) in train_generator:
        img,mask = adjustData(img,mask,flag_multi_class,num_class)       
        yield (img,mask)


def testGenerator(test_path,num_image = 15,target_size = (1280,512),flag_multi_class = False,as_gray = True):
    for i in range(num_image):
        img = io.imread(os.path.join(test_path,"%d.jpg"%i),as_gray = as_gray)
        img = img / 255
        img = trans.resize(img,target_size)
        img = np.reshape(img,img.shape+(1,)) if (not flag_multi_class) else img
        img = np.reshape(img,(1,)+img.shape)
        yield img

def saveResult(save_path,npyfile,flag_multi_class = False,num_class = 2):
    for i,item in enumerate(npyfile):
        img = item[:,:,0]
        io.imsave(os.path.join(save_path,"%d_predict.png"%i),img)



def focal_loss(gamma=2., alpha=0.25):
    def focal_loss_fixed(y_true, y_pred):
        pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
        return -K.sum(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1))-K.sum((1-alpha) * K.pow( pt_0, gamma) * K.log(1. - pt_0))
    return focal_loss_fixed

model=Sequential()
model.add(Conv2D(32,3,strides=1,padding='same',activation='relu',input_shape=(1280,512,1)))
model.add(Conv2D(32,3,strides=1,padding='same',activation='relu'))
model.add(Conv2D(32,3,strides=2,padding='same',activation='relu'))

model.add(Conv2D(64,3,strides=1,padding='same',activation='relu'))
model.add(Conv2D(64,3,strides=1,padding='same',activation='relu'))
model.add(Conv2D(64,3,strides=2,padding='same',activation='relu'))

model.add(Conv2D(128,3,strides=1,padding='same',activation='relu'))
model.add(Conv2D(128,3,strides=1,padding='same',activation='relu'))
model.add(Conv2D(128,3,strides=2,padding='same',activation='relu'))

model.add(UpSampling2D(data_format='channels_last'))
model.add(Conv2D(128,3,padding='same',activation='relu'))
model.add(Conv2D(128,3,padding='same',activation='relu'))

model.add(UpSampling2D(data_format='channels_last'))
model.add(Conv2D(64,3,padding='same',activation='relu'))
model.add(Conv2D(64,3,padding='same',activation='relu'))

model.add(UpSampling2D(data_format='channels_last'))
model.add(Conv2D(32,3,padding='same',activation='relu'))
model.add(Conv2D(32,3,padding='same',activation='relu'))

model.add(Conv2D(1,3,padding='same'))
model.add(Activation(activation='sigmoid'))

data_gen_args = dict(rotation_range=0.2,
                    width_shift_range=0.05,
                    height_shift_range=0.05,
                    shear_range=0.05,
                    zoom_range=0.05,
                    horizontal_flip=True,
                    fill_mode='nearest')
myGene = trainGenerator(2,'F:/py/unet-master/data/mydate/train','image','label',data_gen_args,save_to_dir = None)
model.compile(optimizer = Adam(lr = 1e-4), loss =[focal_loss(alpha=0.25, gamma=2)] , metrics = ['accuracy'])
model.fit_generator(myGene,steps_per_epoch=2000,epochs=1)
model.save('final.h5')
testGene = testGenerator("F:/py/unet-master/data/mydate/test")
model.load_weights("final.h5")
results = model.predict_generator(testGene,5,verbose=1)
saveResult("F:/py/unet-master/data/mydate/test",results)

训练是数据15张，效果如下：

原图：

预测结果：