keras图片多分类

使用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)))

训练验证机准确率很高的情况下，一般来说效果还不错