Python基于ImageAI实现完整的流程：数据集构建、模型训练、识别预测

♚

作者：沂水寒城，CSDN博客专家，个人研究方向：机器学习、深度学习、NLP、CV

Blog: http://yishuihancheng.blog.csdn.net

之前我介绍了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)

基于ImageAI来构建模型和预测真的是代码量少了很多，上面60行左右的代码完成了整理流程的工作，不得不说还是很强大的。

接下来我们开始模型的训练，具体输出如下：

默认只设置了10次迭代，因为迭代次数越多计算时间越长，这里只是一个demo，所以没有迭代很多次，迭代输出结果如下：

训练完成截图如下：

完成了训练的工作后，在我们最初设置的数据集目录data里面会出现类别数据文件和模型数据文件，截图如下：

接下来我们进行简单的预测识别分析，从测试集合里面分别选取一张猫和狗的图像，分别如下：

接下来启动模型预测，输出如下：

从结果上来看，第一张图片识别成了cat，识别错误了，第二张图片识别正确了，这个肯定是由于模型训练所用的数据量太少的问题，一共一个类别才十几张，肯定是不够的，这里只是为了给一个demo，熟悉一下基于ImageAI模块来完整实现真个图片识别的过程。

赞赏作者

Python中文社区作为一个去中心化的全球技术社区，以成为全球20万Python中文开发者的精神部落为愿景，目前覆盖各大主流媒体和协作平台，与阿里、腾讯、百度、微软、亚马逊、开源中国、CSDN等业界知名公司和技术社区建立了广泛的联系，拥有来自十多个国家和地区数万名登记会员，会员来自以公安部、工信部、清华大学、北京大学、北京邮电大学、中国人民银行、中科院、中金、华为、BAT、谷歌、微软等为代表的政府机关、科研单位、金融机构以及海内外知名公司，全平台近20万开发者关注。

▼ 点击成为社区注册会员 「在看」 一下，一起PY！