在前几天的博客里面，我介绍了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)