使用keras实现图片多分类,针对于小图片集效果也还不错。
数据准备,准备两个文件夹 train和validation
文件夹结构为
- train
|--00
|--xx.jpg
|--01
|--xx.jpg
...
训练代码:
# coding=utf8
import tensorflow as tf
import numpy as np
import os
# 基于mobilenet模型
def model_generator():
base_model = tf.keras.applications.mobilenet_v2.MobileNetV2(
include_top=False, weights='imagenet', input_shape=(224, 224, 3))
inputs = tf.keras.layers.Input(shape=(224, 224, 3))
x = base_model(inputs)
# 自己设计顶层并连接
x = tf.keras.layers.GlobalAveragePooling2D()(x)
outputs = tf.keras.layers.Dense(32, activation='softmax')(x)
model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
return model
BS = 32
EPOCHS = 32
train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
rescale=1./255,
shear_range=0.2,
zoom_range=0.2)
validation_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
rescale=1./255,
shear_range=0.2,
zoom_range=0.2)
train_generator = train_datagen.flow_from_directory(
'pics/train',
target_size=(224, 224),
batch_size=BS,
shuffle=True,
class_mode='categorical')
validation_generator = validation_datagen.flow_from_directory(
'pics/validation',
target_size=(224, 224),
batch_size=BS,
shuffle=True,
class_mode='categorical')
# 这步很重要,用于查看class是否跟自己预期一样,或者在上一步 flow_from_directory 中制定class name
print(train_generator.class_indices)
model = model_generator()
model.summary()
model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
loss='categorical_crossentropy',
metrics=['categorical_accuracy'])
if os.path.exists('w.h5'):
model.load_weights('w.h5')
for _ in range(100):
x, y = next(train_generator)
model.fit(x, y, epochs=10, steps_per_epoch=4)
_x, _y = next(validation_generator)
model.evaluate(_x, _y)
model.save_weights('w.h5')
训练图形化显示没做,因为有了mobilenet预训练模型,一般来说训练一会就好了,预测
# coding = utf8
import tensorflow as tf
import cv2
import numpy as np
# 基于mobilenet模型
def model_generator():
base_model = tf.keras.applications.mobilenet_v2.MobileNetV2(
include_top=False, weights='imagenet', input_shape=(224, 224, 3))
inputs = tf.keras.layers.Input(shape=(224, 224, 3))
x = base_model(inputs)
# 自己设计顶层并连接
x = tf.keras.layers.GlobalAveragePooling2D()(x)
outputs = tf.keras.layers.Dense(32, activation='softmax')(x)
model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
return model
model = model_generator()
model.load_weights('w.h5')
for i in range(32):
im = cv2.imread('pics/validation/{:>02}/{:>02}00.jpg'.format(i, i))
im = cv2.resize(im, (224, 224))
im = np.round(im/255, decimals=3)
im = im[..., ::-1] # bgr to rgb
r = model.predict(np.expand_dims(im, axis=0))[0]
print('{}=>{}'.format(i, np.argmax(r)))
训练验证机准确率很高的情况下,一般来说效果还不错
网友评论