Tensorflow+Keras+VGG19 猫狗大战分类

训练程序

import tensorflow as tf
import os
# import pyyaml
import traceback
#读取数据
def read_image_filenames(data_dir):
    cat_dir=data_dir+'cat/'
    dog_dir=data_dir+'dog/'

    #构建特征数据集，值为对应的图片文件名
    cat_filenames=tf.constant([cat_dir + fn for fn in os.listdir(cat_dir)])
    dog_filenames=tf.constant([dog_dir + fn for fn in os.listdir(dog_dir)])
    filenames=tf.concat([cat_filenames,dog_filenames],axis=-1)

    #构建标签数据集，cat为0，dog为1
    labels=tf.concat([
        tf.zeros(cat_filenames.shape,dtype=tf.int32),
        tf.ones(dog_filenames.shape,dtype=tf.int32)],axis=-1)
    return filenames,labels
#解码图片并调整大小
def decode_image_and_resize(filename,label):
    image_string=tf.io.read_file(filename)
    image_decoded=tf.image.decode_jpeg(image_string)
    image_resized=tf.image.resize(image_decoded,[224,224])/255.0
    return image_resized,label


def prepare_dataset(data_dir):
    filenames,labels=read_image_filenames(data_dir)
    dataset=tf.data.Dataset.from_tensor_slices((filenames,labels))
    # 读取图像并处理
    dataset = dataset.map(
        map_func=decode_image_and_resize,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    dataset=dataset.shuffle(buffer_size)
    dataset = dataset.batch(batch_size)
    dataset=dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset


#多少数据进缓冲区，每批次训练多少个
train_data_dir='./dataset/train/'
buffer_size=1300
batch_size=100
dataset_train=prepare_dataset(train_data_dir)
#把Dataset数据集转化为迭代器
it=iter(dataset_train)
images,labels=next(it)

#定义VGG19模型
def vgg19_model(input_shape=(224,224,3)):
    vgg19=tf.keras.applications.vgg19.VGG19(include_top=False,
                                            weights='imagenet',
                                            input_shape=input_shape)
    for layer in vgg19.layers:
        layer.trainable=False
    last=vgg19.output
    #建立新的全连接层
    x=tf.keras.layers.Flatten()(last)
    x=tf.keras.layers.Dense(128,activation='relu')(x)
    x=tf.keras.layers.Dropout(0.3)(x)
    x=tf.keras.layers.Dense(64,activation='relu')(x)
    x=tf.keras.layers.Dropout(0.3)(x)
    x=tf.keras.layers.Dense(128,activation='relu')(x)
    x=tf.keras.layers.Dropout(0.3)(x)
    x=tf.keras.layers.Dense(2,activation='softmax')(x)
    #建立新的模型
    model =tf.keras.models.Model(inputs=vgg19.input,outputs=x)
    # model.summary()
    return model


model=vgg19_model()
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)
#tensorboard
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir="./keras_log/",
                                                      histogram_freq=1,
                                                      )

training_epochs=150
model.fit(dataset_train,epochs=training_epochs,verbose=1,callbacks=[tensorboard_callback])
model.summary()

try:
    #模型保存
    yaml_string=model.to_yaml()
    with open('./models/cat_vs_dog.yaml','w') as model_file:
        model_file.write(yaml_string)
    model.save_weights('./models/cat_vs_dog.h5')
except Exception as e:
    pass

测试程序

import os
import numpy as np
# import cat_vs_dog_train
import tensorflow as tf
# import pyyaml
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing import image

#start为开始读的张数，finish为结束
def read_image_files(path,start,finish,image_size=(224,224)):
    test_file1 = os.listdir(path)
    test_files=[]
    for i in range(len(test_file1)):
        test_files.append(str(i+1)+'.jpg')
    test_image=[]

    #读取图片进行预处理
    for fn in test_files[start:finish]:
        img_filename=path+fn
        img=image.load_img(img_filename,target_size=image_size)
        img_array=image.img_to_array(img)
        test_image.append(img_array)
    test_data=np.array(test_image)
    test_data/=255.0
    print('You choose the image %d to image %d' %(start,finish))
    print("The test_data's shape is",end='')
    print(test_data.shape)
    return test_data
#定义预测函数
def test_image_predict(path,start,finish,image_size=(224,224)):
    #读取测试图片并预处理
    test_data=read_image_files(path,start,finish,image_size)
    #执行预测
    preds=model.predict(test_data)
    #显示图片及预测结果
    for i in range(0,finish-start):
        print(i)
        if np.argmax(preds[i])==0:
            label="cat"+str(preds[i][0])
        else:
            label="dog"+str(preds[i][1])
        plt.title(label)
        plt.imshow(test_data[i])
        plt.axis('off')
        plt.show()

#恢复模型
with open('./models/cat_vs_dog.yaml') as yamlfile:
    loaded_model_yaml=yamlfile.read()
model=tf.keras.models.model_from_yaml(loaded_model_yaml)
#导入权重
model.load_weights('./models/cat_vs_dog.h5')

#执行

test_data_dir='./dataset/test1/'

test_image_predict(test_data_dir,0,1010)