一种基于图像分割实现焊件缺陷检测的方法 | 附源码

import numpy as npimport cv2import osimport randomimport tensorflow as tf
h,w = 512,512
def create_model():
    inputs = tf.keras.layers.Input(shape=(h,w,3))
    conv1 = tf.keras.layers.Conv2D(16,(3,3),activation='relu',padding='same')(inputs)    pool1 = tf.keras.layers.MaxPool2D()(conv1)
    conv2 = tf.keras.layers.Conv2D(32,(3,3),activation='relu',padding='same')(pool1)    pool2 = tf.keras.layers.MaxPool2D()(conv2)
    conv3 = tf.keras.layers.Conv2D(64,(3,3),activation='relu',padding='same')(pool2)    pool3 = tf.keras.layers.MaxPool2D()(conv3)
    conv4 = tf.keras.layers.Conv2D(64,(3,3),activation='relu',padding='same')(pool3)
    upsm5 = tf.keras.layers.UpSampling2D()(conv4)    upad5 = tf.keras.layers.Add()([conv3,upsm5])    conv5 = tf.keras.layers.Conv2D(32,(3,3),activation='relu',padding='same')(upad5)
    upsm6 = tf.keras.layers.UpSampling2D()(conv5)    upad6 = tf.keras.layers.Add()([conv2,upsm6])    conv6 = tf.keras.layers.Conv2D(16,(3,3),activation='relu',padding='same')(upad6)
    upsm7 = tf.keras.layers.UpSampling2D()(conv6)    upad7 = tf.keras.layers.Add()([conv1,upsm7])    conv7 = tf.keras.layers.Conv2D(1,(3,3),activation='relu',padding='same')(upad7)
    model = tf.keras.models.Model(inputs=inputs, outputs=conv7)
    return model
images = []labels = []
files = os.listdir('./dataset/images/')random.shuffle(files)
for f in files:    img = cv2.imread('./dataset/images/' + f)    parts = f.split('_')    label_name = './dataset/labels/' + 'W0002_' + parts[1]    label = cv2.imread(label_name,2)
    img = cv2.resize(img,(w,h))    label = cv2.resize(label,(w,h))
    images.append(img)    labels.append(label)
images = np.array(images)labels = np.array(labels)labels = np.reshape(labels,    (labels.shape[0],labels.shape[1],labels.shape[2],1))
print(images.shape)print(labels.shape)
images = images/255labels = labels/255
model = tf.keras.models.load_model('my_model')
#model = create_model()  # uncomment this to create a new modelprint(model.summary())
model.compile(optimizer='adam', loss='binary_crossentropy',metrics=['accuracy'])model.fit(images,labels,epochs=100,batch_size=10)model.evaluate(images,labels)
model.save('my_model')

import numpy as npimport cv2from google.colab.patches import cv2_imshowimport osimport randomimport tensorflow as tf
h,w = 512,512num_cases = 10
images = []labels = []
files = os.listdir('./dataset/images/')random.shuffle(files)
model = tf.keras.models.load_model('my_model')
lowSevere = 1midSevere = 2highSevere = 4
for f in files[0:num_cases]:    test_img = cv2.imread('./dataset/images/' + f)    resized_img = cv2.resize(test_img,(w,h))    resized_img = resized_img/255    cropped_img = np.reshape(resized_img,          (1,resized_img.shape[0],resized_img.shape[1],resized_img.shape[2]))
    test_out = model.predict(cropped_img)
    test_out = test_out[0,:,:,0]*1000    test_out = np.clip(test_out,0,255)
    resized_test_out = cv2.resize(test_out,(test_img.shape[1],test_img.shape[0]))    resized_test_out = resized_test_out.astype(np.uint16)
    test_img = test_img.astype(np.uint16)
    grey = cv2.cvtColor(test_img, cv2.COLOR_BGR2GRAY)
    for i in range(test_img.shape[0]):     for j in range(test_img.shape[1]):          if(grey[i,j]>150 & resized_test_out[i,j]>40):            test_img[i,j,1]=test_img[i,j,1] + resized_test_out[i,j]            resized_test_out[i,j] = lowSevere          elif(grey[i,j]<100 & resized_test_out[i,j]>40):            test_img[i,j,2]=test_img[i,j,2] + resized_test_out[i,j]            resized_test_out[i,j] = highSevere          elif(resized_test_out[i,j]>40):            test_img[i,j,0]=test_img[i,j,0] + resized_test_out[i,j]            resized_test_out[i,j] = midSevere          else:            resized_test_out[i,j] = 0
    M = cv2.moments(resized_test_out)    maxMomentArea = resized_test_out.shape[1]*resized_test_out.shape[0]*highSevere    print("0th Moment = " , (M["m00"]*100/maxMomentArea), "%")
    test_img = np.clip(test_img,0,255)    test_img = test_img.astype(np.uint8)
    cv2_imshow(test_img)    cv2.waitKey(0)