Keras 二分类 VGG-16

猫狗分类 VGG-16

导入相关的包

from keras.applications.vgg16 import VGG16
from keras.preprocessing import image
from keras.applications.vgg16 import preprocess_input
import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Dropout, Flatten, Dense
from keras.optimizers import SGD
import os

载入模型

可以事先下载好github上的vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5模型，不包括全连接层，然后放在指定目录下。

# 载入预训练的VGG16模型，不包括全连接层 
vgg16_model = VGG16(weights='imagenet', include_top=False, input_shape=(150,150,3))

构建模型

# 搭建全连接层
top_model = Sequential()
top_model.add(Flatten(input_shape=vgg16_model.output_shape[1:]))  #全连接层连接到VGG16的输出 去掉关于批次的维度
top_model.add(Dense(256, activation='relu'))
top_model.add(Dropout(0.5))
top_model.add(Dense(2, activation='softmax'))

model = Sequential()
model.add(vgg16_model)
model.add(top_model)

训练集&测试集数据生成

预先在image/train文件夹下存放若干张猫狗的图片
在image/test文件下存放若干张猫狗的图片

# 训练集数据处理
train_datagen = ImageDataGenerator(
        rescale=1./255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True
)

# 测试集数据处理
test_datagen = ImageDataGenerator(rescale=1./255)

batch_size = 8  #定义批次
# 生成训练数据
train_generator = train_datagen.flow_from_directory(
        'image/train',  # 训练数据路径
        target_size=(150, 150),  # 设置图片大小
        batch_size=batch_size # 批次大小
        ) 

# 测试数据
test_generator = test_datagen.flow_from_directory(
        'image/test',  # 训练数据路径
        target_size=(150, 150),  # 设置图片大小
        batch_size=batch_size # 批次大小
        )

Found 402 images belonging to 2 classes.
Found 273 images belonging to 2 classes.

train_generator.class_indices

{‘cat’: 0, ‘dog’: 1}

训练模型

model.compile(loss='categorical_crossentropy',
              optimizer=SGD(lr=1e-4, momentum=0.9),
              metrics=['accuracy'])

model.fit_generator(
        train_generator,
        steps_per_epoch=len(train_generator),
        epochs=10,
        validation_data=test_generator,
        validation_steps=len(train_generator)
        )
#保存模型
model.save("VGG16.h5")

数据测试

#测试
from keras.models import load_model
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from keras.preprocessing.image import img_to_array,load_img
label = np.array(['cat','dog'])
#载入模型
model = load_model('VGG16.h5')

def deal_img(image):
    image = image.resize((150,150))
    image = img_to_array(image)
    image = image/255
    image = np.expand_dims(image,0)
    return image

for root,dirs,files in os.walk('image/dog_or_cat/'):
    for file in files:
        image_path = os.path.join(root,file)
        print(image_path)
        image_data = load_img(image_path)
        img = Image.open(image_path)
        plt.imshow(img)
        plt.axis('off')
        plt.show()
        image = deal_img(image_data)
        print(label[model.predict_classes(image)])
        print()