基于TensorFlow+OpenCV的焊缝识别
原GitHub项目 跑通后进行了优化和进度显示
资源:
https://download.csdn.net/download/weixin_53403301/85244581
训练部分代码:
import numpy as np
import cv2
import os
import random
import 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/255
labels = labels/255
model = tf.keras.models.load_model('my_model')
model = create_model() # uncomment this to create a new model
print(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 np
import cv2
#import matplotlib.pyplot as plt
#plt.rcParams['font.sans-serif'] = ['SimHei'] # 载入字体
import os
import random
import tensorflow as tf
h,w = 512,512
num_cases = 3
images = []
labels = []
files = os.listdir('./dataset/images/')
random.shuffle(files)
model = tf.keras.models.load_model('my_model')
lowSevere = 1
midSevere = 2
highSevere = 4
#def pshow(words,picture):
# plt.imshow(picture[:,:,::-1])
# plt.title(words), plt.xticks([]), plt.yticks([])
# plt.show()
for f in files[0:num_cases]:
print(0,f)
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]))
print(1,f)
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)
print("LOOP:",f)
img_len = test_img.shape[0]
print(img_len)
for i in range(img_len):
print("LOOP:",f,str(int(i/img_len*100))+"%")
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
print("END_LOOP:",f)
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)
# pshow(str(f),test_img)
# cv2.imshow(str(f),test_img)
cv2.imwrite('./save/end_'+f,test_img)
# cv2.waitKey(0)