model.rar

import tensorflow as tf import numpy as np from tensorflow.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array from tensorflow.keras.layers import Dense, Flatten, Conv2D, Dropout, MaxPooling2D from tensorflow.keras.optimizers import Adam from tensorflow.keras.losses import sparse_categorical_crossentropy from tensorflow.keras.callbacks import ModelCheckpoint import cv2 class Mask(object): def __init__(self): # 定义训练和测试图片的变化方法，标准化及数据增强 self.train_generator = ImageDataGenerator(rescale=1.0 / 255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True) self.test_generator = ImageDataGenerator(rescale=1.0 / 255) # 指定训练集合测试集的目录 self.train_dir = './data/train' self.test_dir = './data/test' # 定义图片训练相关的网络参数 self.image_size = (150, 150) self.batch_size = 32 # 每次返回数据的个数 self.label_dict = { '0': 'mask', '1': 'unmask'} def get_local_data(self): """ 读取本地图片数据 :return:训练数据和测试数据的迭代器 """ # 使用flow_from_directory train_gen = self.train_generator.flow_from_directory(self.train_dir, target_size=self.image_size, batch_size=self.batch_size, class_mode="binary", shuffle=True) test_gen = self.test_generator.flow_from_directory(self.test_dir, target_size=self.image_size, batch_size=self.batch_size, class_mode="binary", shuffle=True) return train_gen, test_gen def mask_model(self): # 口罩检测模型 model = tf.keras.models.Sequential([ # 1 2维卷积层 Conv2D(100, (3, 3), activation='relu', input_shape=(150, 150, 3)), # 2 池化层 tf.keras.layers.MaxPooling2D(2, 2), # 3 2维卷积层 Conv2D(100, (3, 3), activation='relu'), # 4 池化层 MaxPooling2D(2, 2), # 5 Flatten层 Flatten(), # 6 Dropout层 Dropout(0.5), # 7 Dense又被称作全连接层 Dense(50, activation='relu'), # 8 Dense层 Dense(2, activation='softmax') ]) return model def compile(self, model): """ 优化、损失、accuracy :return: None """ model.compile(optimizer=Adam(), loss=sparse_categorical_crossentropy, metrics=['accuracy']) return None def fit_generator(self, model, train_gen, test_gen): """ 训练模型， model.fit_generator()而不是选择model.fit（） :param model: :return: """ checkpoint = ModelCheckpoint('./ckpt/mask_{epoch:02d}-{val_loss:.2f}.h5', monitor='val_loss', verbose=0, save_best_only=True, mode='auto') model.fit(train_gen, validation_data=test_gen, epochs=3, callbacks=[checkpoint]) return None def frame(self): webcam = cv2.VideoCapture(0) def predict_1(self, model, frame): model.load_weights('./mask_model/ckpt/mask_03-0.13.h5') normalized = frame / 255.0 reshaped = np.reshape(normalized, (1, 150, 150, 3)) reshaped = np.vstack([reshaped]) result = model.predict(reshaped) max_probability = result[0][0] return max_probability if __name__ == '__main__': tm = Mask() # train_gen, test_gen = tm.get_local_data() model = tm.mask_model() # tm.compile(model) # tm.fit_generator(model, train_gen, test_gen) tm.predict_1(model)