0031-keras入门-图片多分类

布满荆棘的人生 2023-01-11 01:29 131阅读 0赞

from keras.applications.resnet50 import ResNet50
    from keras.applications.vgg16 import VGG16
    from keras.models import Model
    from keras.layers import Dropout,Flatten,Dense
    from keras.optimizers import SGD,Adam
    from keras.callbacks import TensorBoard
    from keras.preprocessing.image import ImageDataGenerator
    from keras import backend
    from keras.callbacks import ModelCheckpoint
    from matplotlib import pyplot as plt
    import matplotlib as mpl
    import tensorflow as tf
    from keras import backend as K
    from PIL import ImageFile
    import os
    ImageFile.LOAD_TRUNCATED_IMAGES = True
    mpl.use('TkAgg')
    
    #进行配置
    os.environ["CUDA_VISIBLE_DEVICES"] = "0" #使用编号为0，1的GPU
    
    class Train():
        def __init__(self,model_type):
            if model_type=='resnet50':
                self.target_size = (224,224) #图片将被resize成该尺寸
                self.weights_dir = 'model/resnet50_weights.h5'
            elif model_type=='vgg16':
                self.target_size = (224,224)
            self.train_dir = '/opt/zhongls/object_detect/1145_old/train'  # 训练数据保存路径
            self.val_dir = '/opt/zhongls/object_detect/1145_old/val'  # 验证数据保存路径
            self.test_dir = '/opt/zhongls/object_detect/1145_old/test'  # 测试数据保存路径
            self.num = 1145  # 多少分类
    
        #图片读取和图片扩增和归一化处理
        def unit_img(self):
            train_datagen = ImageDataGenerator(rescale=1./255,
                                               shear_range=0.2,
                                               zoom_range=0.2,
                                               horizontal_flip=True
                                               )
            test_datagen = ImageDataGenerator(rescale=1./255)
            val_datagen = ImageDataGenerator(rescale=1./255)
            # 图片归一化处理
            train_flow = train_datagen.flow_from_directory(self.train_dir,  # 训练目录
                                                           self.target_size,  # 所有图像调整为
                                                           batch_size=32,
                                                           class_mode='categorical')
            #使用的损失函数，多分类设为categorical
            val_flow = val_datagen.flow_from_directory(self.val_dir,self.target_size,batch_size=32,class_mode='categorical')
            test_flow = test_datagen.flow_from_directory(self.test_dir, self.target_size, batch_size=32,
                                                         class_mode='categorical')
            return train_flow, val_flow, test_flow
        
        def creat_model(self,model_type):
            if model_type == 'resnet50':
                base_model = ResNet50(include_top=True, weights='imagenet')
                print(len(base_model.layers))
                base_model = Model(input=base_model.input, output=base_model.get_layer(index=175).output)
                x = base_model.output
                x = Dense(250,activation='relu')(x)
                predictions = Dense(self.num, activation='softmax')(x)
                model = Model(inputs=base_model.input, outputs=predictions)
                model.summary()
            elif model_type == 'vgg16':
                base_model = VGG16(include_top=True, weights=None)
                print(len(base_model.layers))
                base_model = Model(input=base_model.input, output=base_model.get_layer(index=21).output)
                x = base_model.output
                predictions = Dense(self.num, activation='softmax')(x)
                model = Model(inputs=base_model.input, outputs=predictions)
                model.summary()
            return model
    
        #绘制训练&验证的损失值
        def show(self,history):
            plt.plot(history.history['accuracy'])
            plt.plot(history.history['val_accuracy'])
            plt.title('Model accuracy')
            plt.ylabel('Accuracy')
            plt.xlabel('Epoch')
            plt.legend(['Train', 'Test'], loc='upper left')
            plt.show()
    
            plt.plot(history.history['loss'])
            plt.plot(history.history['val_loss'])
            plt.title('Model loss')
            plt.ylabel('Loss')
            plt.xlabel('Epoch')
            plt.legend(['Train', 'Test'], loc='upper left')
            #plt.show()
            plt.savefig('2.jpg')
    
        #编译和训练模型
        def train(self,model_type):
            model = self.create_model(model_type=model_type)
            model.compile(optimizer=Adam(lr=0.0001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0005),
                          loss='categorical_crossentropy',
                          metrics=['accuracy'])
            train_flow,val_flow,test_flow = self.unit_img()
    
            #训练
            save_dir = '/opt/zhongls/object_detect/1145_old/model/'
            checkpointer = ModelCheckpoint(os.path.join(save_dir, 'model_{epoch:05d}_{loss:.5f}.hdf5'),
                                           verbose=1, save_weights_only=False, period=100)
    
            # history = model.fit_generator(datagenerator(batch_size=xxx),
            #                               steps_per_epoch=xxx, epochs=xxx, verbose=1,
            #                               initial_epoch=initial_epoch,
            #                               callbacks=[checkpointer])
            history = model.fit_generator(train_flow,
                                          steps_per_epoch=500,
                                          epochs=10000,
                                          validation_data=val_flow,
                                          validation_steps=12,
                                          initial_epoch=0,
                                          callbacks=[checkpointer])
            test_loss, test_acc = model.evaluate_generator(test_flow, 96)
            print("test acc:",test_acc)
            #保存模型和权重
            model.save(self.weights_dir)
            self.show(history=history)
    
        if __name__ == '__main__':
            config = tf.ConfigProto()
            config.gpu_options.allow_growth = True
            sess = tf.Session(config=config)
            K.set_session(sess)
            print("网络结构：vgg16,resnet50")
            model_type = 'resnet50'
            tr = Train(model_type=model_type)
            tr.train(model_type=model_type)