Keras机器视觉小tip——以垃圾分类为例

垃圾分类的目标是，对于给出的图片进行分类，所以需要一个分类网络。

一。首先了解一下数据集

点开dataset-resized看一下，这里将所有的图片分成六个类别：

除了图片的放置与命名，再来看一下标注文件：

 

TIP 1：介绍一下python中的glob库

glob是python自己带的一个文件操作相关模块，用它可以查找符合自己目的的文件，就类似于Windows下的文件搜索，支持通配符操作，*,?,[]这三个通配符，*代表0个或多个字符，?代表一个字符，[]匹配指定范围内的字符，如[0-9]匹配数字。

import glob,os

data_path = "D:\pythonTest\datasets\la1ji1fe1nle4ishu4ju4ji22-momodel\dataset-resized"

img_list = glob.glob(os.path.join(data_path,'*/*.jpg'))
print(len(img_list))

这样就可以获取到图片的数量：

TIP 2 :在pycharm中使用plt.imshow() 有可能会显示不出来，可以导入pylab包。

现在我们随机显示一下数据集中的6张图片：

import random
import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import pylab

for i,img_path in enumerate(random.sample(img_list,6)):
    img = cv.imread(img_path)
    img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    plt.subplot(2,3,i+1)
    plt.imshow(img)
    plt.axis('off')
pylab.show()

这里的enumerate()函数会返回两个值，下标和内容。

再获取一下图片的尺寸：

path = random.sample(img_list,1)
img = cv.imread(path[0])
print(img.shape)

得到的结果是（384，512，3）

 

二。对数据集有了初步的了解后，需要对数据集进行预处理

这里用一个keras的模块，详情参考https://keras-cn.readthedocs.io/en/latest/preprocessing/image/

from keras.preprocessing.image import ImageDataGenerator

def processing_data(data_path,height,width,batch_size =32,validation_split = 0.1):
    train_data = ImageDataGenerator(
        rescale=1./255,
        shear_range=0.1,
        zoom_range=0.1,
        width_shift_range=0.1,
        height_shift_range=0.1,
        horizontal_flip=True,
        vertical_flip=True,
        validation_split=validation_split
    )

    validation_data = ImageDataGenerator(
        rescale=1./255,
        validation_split = validation_split
    )

    train_generator = train_data.flow_from_directory(
        data_path,
        target_size=(height,width),
        batch_size = batch_size,
        class_mode = 'categorical',
        subset='training',
        seed = 0
    )

    validation_generator = validation_data.flow_from_directory(
        data_path,
        target_size=(height, width),
        batch_size=batch_size,
        class_mode='categorical',
        subset='validation',
        seed=0
    )

    return train_generator,validation_generator

data_path = "D:\pythonTest\datasets\la1ji1fe1nle4ishu4ju4ji22-momodel\dataset-resized"
height,width = 384,512
train_generator,validation_generator = processing_data(data_path,height,width)
labels = train_generator.class_indices
print(labels)
labels = dict((v,k)for k,v in labels.items())
print(labels)

class_indices可以获得标签和对应的类别序号

3。创建模型

#创建模型
def dnn_model(input_shape,train_generator,validation_generator,model_save_path='results/try.h5',log_dir = "results/logs/"):
    resize = 224
    classes = 6
    prob = 0.5

    model = Sequential()
    model.add(Conv2D(64,(3,3),strides=(1,1),input_shape = input_shape,padding='same',activation='relu',kernel_initializer='uniform'))
    model.add(Conv2D(64, (3, 3), strides=(1, 1), padding='same', activation='relu',kernel_initializer='uniform'))
    model.add(MaxPooling2D(pool_size=(2,2)))
    model.add(Conv2D(128, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(Conv2D(128, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(256, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(Conv2D(256, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(Conv2D(256, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(512, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(Conv2D(512, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(Conv2D(512, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(512, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(Conv2D(512, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(Conv2D(512, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Flatten())
    model.add(Dense(4096,activation='relu'))
    model.add(Dropout(prob))
    model.add(Dense(4096,activation='relu'))
    model.add(Dropout(prob))
    model.add(Dense(classes,activation='softmax'))

这里以vgg16网络为例。经过几轮的卷积和池化，最后Flatten，用了两个全连接和Dropout，最后使用分类函数softmax输出六个分类。这一部分可以自行改变，尝试使用其他网络结构。

4。模型编译

# 编译模型, 采用 compile 函数: https://keras.io/models/model/#compile
    model.compile(
            # 是优化器, 主要有Adam、sgd、rmsprop等方式。
            optimizer='Adam',
            # 损失函数,多分类采用 categorical_crossentropy
            loss='categorical_crossentropy',
            # 是除了损失函数值之外的特定指标, 分类问题一般都是准确率
            metrics=['accuracy'])

    # 可视化，TensorBoard 是由 Tensorflow 提供的一个可视化工具。
    tensorboard = TensorBoard(log_dir)

关于优化器的区别和选择：https://blog.csdn.net/brucewong0516/article/details/78838124

5。模型训练

 # 训练模型, fit_generator函数:https://keras.io/models/model/#fit_generator
    # 利用Python的生成器，逐个生成数据的batch并进行训练。
    # callbacks: 实例列表。在训练时调用的一系列回调。详见 https://keras.io/callbacks/。
    d = model.fit_generator(
            # 一个生成器或 Sequence 对象的实例
            generator=train_generator,
            # epochs: 整数，数据的迭代总轮数。
            epochs=5,
            # 一个epoch包含的步数,通常应该等于你的数据集的样本数量除以批量大小。
            steps_per_epoch=2076 // 32,
            # 验证集
            validation_data=validation_generator,
            # 在验证集上,一个epoch包含的步数,通常应该等于你的数据集的样本数量除以批量大小。
            validation_steps=231 // 32,
            callbacks=[tensorboard])
    # 模型保存
    model.save(model_save_path)

6。模型加载与评估

#加载并评估模型
def load_and_model_prediction(validation_generator,model_save_path):
    model = load_model(model_save_path)
    loss,accuracy = model.evaluate_generator(validation_generator)
    print("\nloss:%.2f,Accuracy:%.2f%%"%(loss,accuracy*100))

7。运行上面的函数

def main():
    data_path = "D:\pythonTest\datasets\la1ji1fe1nle4ishu4ju4ji22-momodel\dataset-resized"
    save_model_path = "results/try.h5"

    #获取训练集和测试集
    train_data,test_data = processing_data(data_path)
    input_shape = (384,512,3)

    #进行模型训练
    model(input_shape,train_data,test_data,save_model_path)

    #模型苹果
    load_and_model_prediction(test_data,save_model_path)

if __name__ == '__main__':
    main()

运行vgg16直接把我的pycharm卡退出，所以尽量用服务器GPU跑吧。

最后，介绍一下当我们得到模型文件后（h5文件），怎么来使用这个模型：

import os
from keras.models import load_model
from keras.preprocessing import image
import numpy as np


def load_and_predict(img):
    """
    加载模型并预测一张图片的类别
    :param img: PIL.Image 对象
    :return: string, 模型识别图片的类别, 
            共 'cardboard','glass','metal','paper','plastic','trash' 6 个类别
    """
    # 加载模型, 默认'results/dnn.h5',请填写你的最佳模型
    model_path = 'results/dnn.h5'
    model = load_model(model_path)

    # 把图片转换成为numpy数组
    img = image.img_to_array(img)
    
    # 图片放缩
    img = 1.0/255 * img

    # expand_dims的作用是把img.shape转换成(1, img.shape[0], img.shape[1], img.shape[2])
    x = np.expand_dims(img, axis=0)

    # 模型预测
    y = model.predict(x)

    # 获取labels
    labels = {0: 'cardboard', 1: 'glass', 2: 'metal', 3: 'paper', 4: 'plastic', 5: 'trash'}

    # 获取输入图片的类别
    y_predict = labels[np.argmax(y)]

    # 返回图片的类别
    return y_predict

if __name__ == '__main__':
    path = "D:/pythonTest/img/paper9.jpg"
    img = image.load_img(path)
    print(load_and_predict(img))