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