ImageAI實現完整的流程：資料集構建、模型訓練、識別預測

AIBigbull2050發表於2019-08-28

在前幾天的部落格裡面，我介紹了ImageAI模組，以及基於官方預訓練得到的模型進行了簡單的識別分析，整體的流程都是比較簡單的，今天想要基於ImageAI模組來完整地實現整個圖片分類識別的流程，也就是說：這裡的資料集構建和模型的訓練以及結果模型的呼叫預測都是自己完成的。ImageAI簡化了模型的搭建流程，所以整體來說還是比較簡單的。

按照官方的講解我們先來構建自己本地的資料集，我這裡採用的是之前kaggle比賽中使用到的貓狗大戰資料集，由於PC機計算能力有限，這裡的資料量我只使用了很少的一部分。

首先建立一個資料夾data，在data裡面建立兩個資料夾，分別是：train和test，看名字就是分別存放的是訓練資料集和測試資料集，在在train和test裡面分別建立各個不同類別的資料資料夾，每個類別單獨存放於一個資料夾中，我們從下面的截圖來簡單看看資料集整體的結構，如下所示：

進入data如下：

進入train如下：

進入cat如下：

完成了資料集的構建後，我們就要開始訓練了，具體實現如下：

#!usr/bin/env python

# encoding:utf-8

from __future__ import division

 

"""

__Author__:沂水寒城

功能： Python基於ImageAI實現完成的流程：資料集構建、模型訓練、識別預測

"""

 

 

import os

import threading

from imageai.Prediction import ImagePrediction

from imageai.Prediction.Custom import ModelTraining

from imageai.Prediction.Custom import CustomImagePrediction

 

 

 

def modelTrain(dataDir='data',classNum=2,epochs=100,batch_size=32):

 '''

 模型訓練部分

 '''

 #建立了ModelTraining類的新例項

 model_trainer = ModelTraining()

 #將模型型別設定為ResNet

 model_trainer.setModelTypeAsResNet()

 #設定我們想要訓練的資料集的路徑

 model_trainer.setDataDirectory(dataDir)

 #模型訓練

 '''

 num_objects：該引數用於指定影像資料集中物件的數量

 num_experiments：該引數用於指定將對影像訓練的次數，也稱為epochs

 enhance_data（可選）：該引數用於指定是否生成訓練影像的副本以獲得更好的效能。

 batch_size：該引數用於指定批次數量。由於記憶體限制，需要分批訓練，直到所有批次訓練集都完成為止。

 show_network_summary：該引數用於指定是否在控制檯中顯示訓練的過程。

 '''

 model_trainer.trainModel(num_objects=classNum, num_experiments=epochs,enhance_data=True, 

 batch_size=batch_size, show_network_summary=True)

 print('Model Train Finished!!!')

 

 

 

def modelPredict(model_path='data/models/model_ex-001_acc-0.500000.h5',

 class_path='data/json/model_class.json',

 pic_path='a.jpg',classNum=2,resNum=5):

 '''

 模型預測部分

 prediction_speed[模型載入的速度]: fast faster fastest

 '''

 prediction=CustomImagePrediction()

 prediction.setModelTypeAsResNet()

 prediction.setModelPath(model_path)

 prediction.setJsonPath(class_path)

 prediction.loadModel(num_objects=classNum,prediction_speed='fastest') 

 predictions,probabilities=prediction.predictImage(pic_path,result_count=resNum)

 for eachPrediction, eachProbability in zip(predictions, probabilities):

 print(eachPrediction+" : "+str(eachProbability))

 

 

 

if __name__=='__main__':

 #模型訓練

 modelTrain(dataDir='data',classNum=2,epochs=10,batch_size=8)

 

 #模型識別預測

 modelPredict(model_path='data/models/model_ex-001_acc-0.500000.h5',

 class_path='data/json/model_class.json',

 pic_path='test.jpg',classNum=2,resNum=5)