ARJ_DenseNet BMR模型训练

废话不多数，模型训练代码

densenet_arj_BMR.py

：

import timefrom tensorflow.keras.applications.xception import Xception
from tensorflow.keras.applications.densenet import DenseNet169
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow.keras as keras
from arj_t.plt_graph import show_graph
from common_para import train_dir, val_dir, station, EPOCHS_1,EPOCHS_2, batch_size, CLASS_WEIGHT, classesinput_shape = (224, 224)
date_ = time.strftime('%Y%m%d', time.localtime())
cpkt_path = f'./ckpt/ARJ_Densenet_ckpt{station}_20231017-1.h5'
model_path = f'./ckpt/ARJ_Densenet_MODEL{station}_{date_}.h5'class ArjDensenetModel(object):def __init__(self):self.base_model = DenseNet169(weights='imagenet', include_top=False)# 泛化能力不行，进行图像增强测试self.train_gen = ImageDataGenerator(rescale=1.0 / 255.0,# rotation_range=45,# width_shift_range=0.2,# height_shift_range=0.2,# brightness_range=(0, 0.3),# shear_range=0.2, # 浮点数。剪切强度（以弧度逆时针方向剪切角度）# zoom_range=[0.5, 1.5],  # 小于1.0的缩放将放大图像，大于1.0的缩放将缩小图像。# horizontal_flip=True,# vertical_flip=True,# fill_mode='constant',# cval=0)# self.train_gen = ImageDataGenerator(rescale=1.0 / 255.0)self.val_gen = ImageDataGenerator(rescale=1.0 / 255.0)# 获取本地训练和验证图片，生成generatordef get_local_data(self):self.train_gen = self.train_gen.flow_from_directory(directory=train_dir,target_size=input_shape,batch_size=batch_size,class_mode='binary',  # binary 改为 categoricalshuffle=True,# save_to_dir=r'D:\AOI Gray Image-OA\dataset\BMR\train_trans2',# save_format='jpg',# save_prefix='trans_')self.val_gen = self.val_gen.flow_from_directory(directory=val_dir,target_size=input_shape,batch_size=batch_size,class_mode='binary',  # binary 改为 categorical 2022/5/15shuffle=True)return Nonedef refine_basemode(self):"""获取VGG16 basemode只获取全连接层以前的卷积和池化层，并进行参数冻结，也就是使用原有训练好的参数自主增加隐藏层和全连接层进行训练，获得目标模型:return:"""# 获取除全连接层以外的层数,no-top modelx = self.base_model.outputs[0]# 加入全局池化、隐藏层、全连接层x = keras.layers.GlobalAveragePooling2D()(x)x = keras.layers.Dense(2048, activation='relu')(x)# x = keras.layers.BatchNormalization()(x)x = keras.layers.Dense(1024, activation='relu')(x)out = keras.layers.Dense(2, activation='softmax')(x)# 生成新的模型new_model = keras.models.Model(inputs=self.base_model.inputs, outputs=out)# 冻结vgg模型原有参数self.freeze_base_model()# 对new_model进行编译# 学习效果不佳，初始学习率加大尝试# 初始学习率0.01->0.001opt = keras.optimizers.Adam(learning_rate=0.001)new_model.compile(# optimizer=opt,  # 优化器# # 因为class_mode使用了categorical, 此时返回one-hot编码标签# # 那么这里就需要使用categorical_crossentropy，多类对数交叉熵损失计算# # 如果class_mode使用binary, 此时返回1D的二值标签，loss就需要使用sparse_categorical_crossentropy# loss='sparse_categorical_crossentropy',  # 使用交叉熵损失函数 分类# metrics=['accuracy']# binary_crossentropy与sigmoid联合使用二分类# categorical_crossentropy与softmax联合使用optimizer='adam',loss='sparse_categorical_crossentropy',metrics=['accuracy'])return new_model# 冻结模型训练层数def freeze_base_model(self):for layer in self.base_model.layers:layer.trainable = Falsereturn None# # 对new_model进行trainingdef fit(self, model):# 获取本地数据self.get_local_data()# 定义checkpointckpt = keras.callbacks.ModelCheckpoint(filepath=cpkt_path,monitor='val_accuracy',save_freq='epoch',save_weights_only=True,save_best_only=True)# 早停法用起来el1 = keras.callbacks.EarlyStopping(monitor='val_accuracy',patience=15,verbose=2,mode='auto')# 定义学习率缩小规则rc1 = keras.callbacks.ReduceLROnPlateau(monitor='val_accuracy',factor=0.1,  # 学习率缩小倍数 new_lr = lr*factorpatience=5,  # 耐心吗，5次迭代不增加就缩小学习率mode='auto',verbose=1,  # 1代表更新信息，0代表不更新# epsilon=0.0001,  # 确认是否进入平原区min_lr=0,cooldown=0)# 模型训练# 加入class_weight权重# 暂时注释。his1 = model.fit(self.train_gen, validation_data=self.val_gen,epochs=EPOCHS_1, callbacks=[ckpt, rc1, el1])# his1 = model.fit(self.train_gen, validation_data=self.val_gen,#                  epochs=EPOCHS_1, callbacks=[ckpt, rc1, el1], class_weight=CLASS_WEIGHT)print('first step end')# 解冻所有layer，进行参数微调for layer in model.layers:layer.trainable = True# 早停法用起来el2 = keras.callbacks.EarlyStopping(monitor='val_accuracy',patience=11,verbose=2,mode='auto')# 定义学习率缩小规则rc2 = keras.callbacks.ReduceLROnPlateau(monitor='val_accuracy',factor=0.1,  # 学习率缩小倍数 new_lr = lr*factorpatience=5,  # 耐心吗，5次迭代不增加就缩小学习率mode='auto',verbose=1,  # 1代表更新信息，0代表不更新# epsilon=0.0001,  # 确认是否进入平原区min_lr=0,cooldown=0)opt = keras.optimizers.Adam(learning_rate=0.001)model.compile(optimizer=opt,loss='sparse_categorical_crossentropy',metrics=['accuracy'])# 模型训练# model.load_weights(cpkt_path)his2 = model.fit(self.train_gen, validation_data=self.val_gen,epochs=EPOCHS_2, callbacks=[ckpt, rc2, el2], class_weight={0: 1, 1: 1.5})# # 模型训练# his2 = model.fit(self.train_gen, validation_data=self.val_gen,#                  epochs=EPOCHS_2, callbacks=[ckpt, rc2, el2], class_weight={0: 1, 1: 2, 2: 3})print('END STEP')return his1, his2if __name__ == '__main__':arj_model = ArjDensenetModel()model = arj_model.refine_basemode()his1, his2 = arj_model.fit(model)# # 保存模型# model.save(model_path)show_graph(his1)show_graph(his2)

common_para.py代码

train_dir = r"D:\new_data\BMR_TRAIN\train"
val_dir = r"D:\new_data\BMR_TRAIN\validate"
station = '_ALL_BMR'
batch_size = 32
EPOCHS_1 = 10
EPOCHS_2 = 40
CLASS_WEIGHT = {0: 1., 1: 1., 2: 1.}
threshold_value = 0
classes = 2

模型预测代码 

BMR_IPS_135K_predict.py

import osimport numpy as np
from tensorflow.keras.preprocessing.image import load_img, img_to_arrayimport densenet_arj_BMR
import inceptionRestnet_arj_t
import resnet101_arj_BMR
import xception_arj_BMRMODEL_NAME = 'densenet'val_path = r'D:\AOI Gray Image-OA\dataset\case1\135K-ISR-IPS\validate'
# val_path = r'D:\AOI Gray Image-OA\dataset\error\W_to_G'
other_path = r'D:\AOI Gray Image-OA\AOI IMAGE-20220513\A6Q\ISR\复判后-G'
test_path = r'D:\new_data\BMR表外测试\BMR\A1A\P'
# TARGET_SIZE = (299, 299)
# DEFECT_TYPE = 'P'
# error_path = fr"D:\AOI Gray Image-OA\dataset\BMR\{MODEL_NAME}"def get_ckptpath_model():# arj_model = tensorflow.keras.models.Model()ckpt_path = ''# target_size = (224, 224)if MODEL_NAME == 'xception':ckpt_path = xception_arj_BMR.cpkt_patharj_model = xception_arj_BMR.ArjResnet101Model()target_size = xception_arj_BMR.input_shapeelif MODEL_NAME == 'inceptionRestnet':ckpt_path = inceptionRestnet_arj_t.cpkt_patharj_model = inceptionRestnet_arj_t.ArjInceptionRestnetModel()target_size = inceptionRestnet_arj_t.input_shapeelif MODEL_NAME == 'densenet':ckpt_path = densenet_arj_BMR.cpkt_patharj_model = densenet_arj_BMR.ArjDensenetModel()target_size = densenet_arj_BMR.input_shapeelif MODEL_NAME == 'resnet101':ckpt_path = resnet101_arj_BMR.cpkt_patharj_model = resnet101_arj_BMR.ArjResnet101Model()target_size = resnet101_arj_BMR.input_shapereturn ckpt_path, arj_model, target_size# 获取想要预测的图片绝对路径，包含文件名
def get_img_paths(defect_type, path):img_path = os.path.join(path, defect_type)img_paths = []for root, dirs, files in os.walk(img_path):for file in files:# print(file[-3:])if file[-3:] == 'jpg':img_paths.append(os.path.join(root, file))return img_pathsdef bmr_ips_predict(img_paths, error_path, defect_type='G'):ckpt_path, arj_model, input_shape = get_ckptpath_model()model = arj_model.refine_basemode()print(ckpt_path)model.load_weights(ckpt_path)print(model.summary())predict_dict = {0: 'G', 1: 'P', 2: 'W'}# 加载图片，预测white_cnt = 0good_cnt = 0repair_cnt = 0threshold_ls = []for img_path in img_paths:img_arr = load_img(img_path, target_size=input_shape)img = img_arr# print(img_path)# 转化为矩阵img_arr = img_to_array(img_arr)# print(img.shape)# 归一化# img_arr = preprocess_input(img_arr)img_arr /= 255.# print(type(img_arr))# img_arr = preprocess_input(img_arr)# img_arr /= 127.5# img_arr -= 1.# 形状修改img_arr = img_arr.reshape(1, img_arr.shape[0], img_arr.shape[1], img_arr.shape[2])# print(img.shape)# print(img_arr)y_predict = model.predict(img_arr)index = np.argmax(y_predict)# 加入阈值threshold = y_predict[0][index]# print(img_path.split('\\')[-1])# print(y_predict[0], ' >> ', threshold)# threshold_ls.append(threshold)# print(y_predict)y_predict = predict_dict[index]# print(index)# print(y_predict)# if index == 0:#     good_cnt += 1# else:#     repair_cnt += 1# 保存判错的图片# 预测结果G# save_img_name = str(round(threshold,2))+'_'+img_path.split('\\')[-1]save_img_name = img_path.split('\\')[-1]if index == 0:# 加入阈值调节判G能力if threshold > 0:good_cnt += 1# print(good_cnt)# print(img_path[-10:])# 如果原本P文件夹if defect_type == 'P':threshold_ls.append(threshold)img.save(os.path.join(error_path, 'AI_P_TO_G', save_img_name))os.remove(img_path)# 如果原本W文件夹if defect_type == 'W':img.save(os.path.join(error_path, 'AI_W_TO_G', save_img_name))os.remove(img_path)else:repair_cnt += 1elif index == 1:repair_cnt += 1if defect_type == 'G':threshold_ls.append(threshold)img.save(os.path.join(error_path, 'AI_G_TO_P', save_img_name))os.remove(img_path)if defect_type == 'W':img.save(os.path.join(error_path, 'AI_W_TO_P', save_img_name))os.remove(img_path)elif index == 2:white_cnt += 1if defect_type == 'G':img.save(os.path.join(error_path, 'AI_G_TO_W', save_img_name))os.remove(img_path)if defect_type == 'P':img.save(os.path.join(error_path, 'AI_P_TO_W', save_img_name))os.remove(img_path)else:print('还有第四种可能？？！！')# print(y_predict)# print('**************************')# pd.DataFrame(data=threshold_ls).to_csv('./threshold.csv', encoding='utf-8')print(threshold_ls)print('good_cnt :  %d' % good_cnt)print('repair_cnt :  %d' % repair_cnt)# print('white_cnt :  %d' % white_cnt)# if __name__ == '__main__':
#     paths = get_img_paths(DEFECT_TYPE, test_path)
#     bmr_ips_predict(paths,error_path)

模型总预测代码

predict_all.py

import BMR_IPS_135K_predict# 此程序用来进行所有模型预测2023/10/17img_path = r'D:\new_data\BMR_TRAIN\test\WHITE'DEFECT_TYPE = 'P'paths = BMR_IPS_135K_predict.get_img_paths(DEFECT_TYPE, img_path)BMR_IPS_135K_predict.bmr_ips_predict(paths, img_path, DEFECT_TYPE)