【AI实战】动手训练自己的目标检测模型(YOLO篇)

原创
2018/08/14 18:49
阅读数 29.1W

在前面的文章中,已经介绍了基于SSD使用自己的数据训练目标检测模型(见文章:手把手教你训练自己的目标检测模型),本文将基于另一个目标检测模型YOLO,介绍如何使用自己的数据进行训练。

 
YOLO(You only look once)是目前流行的目标检测模型之一,目前最新已经发展到V3版本了,在业界的应用也很广泛。YOLO的基本原理是:首先对输入图像划分成7x7的网格,对每个网格预测2个边框,然后根据阈值去除可能性比较低的目标窗口,最后再使用边框合并的方式去除冗余窗口,得出检测结果,如下图:
 
YOLO的特点就是“快”,但由于YOLO对每个网格只预测一个物体,就容易造成漏检,对物体的尺度相对比较敏感,对于尺度变化较大的物体泛化能力较差。

本文的目标仍旧是在图像中识别检测出可爱的熊猫

基于YOLO使用自己的数据训练目标检测模型,训练过程跟前面文章所介绍的基于SSD训练模型一样,主要步骤如下:
 
1、安装标注工具
本案例采用的标注工具是labelImg,在前面的文章介绍训练SSD模型时有详细介绍了安装方法(见文章:手把手教你训练自己的目标检测模型),在此就不再赘述了。
成功安装后的labelImg标注工具,如下图:

2、标注数据
使用labelImg工具对熊猫照片进行画框标注,自动生成VOC_2007格式的xml文件,保存为训练数据集。操作方式跟前面的文章介绍训练SSD模型的标注方法一样(见文章:手把手教你训练自己的目标检测模型),在此就不再赘述了。

3、配置YOLO
(1)安装Keras
 
本案例选用YOLO的最新V3版本,基于Keras版本。Keras是一个高层神经网络API,以Tensorflow、Theano和CNTK作为后端。由于本案例的基础环境(见文章:AI基础环境搭建)已经安装了tensorflow,因此,Keras底层将会调用tensorflow跑模型。Keras安装方式如下:

# 切换虚拟环境
source activate tensorflow
# 安装keras-gpu版本
conda install keras-gpu
# 如果是安装 keras cpu版本,则执行以下指令
#conda install keras

keras版本的yolo3还依赖于PIL工具包,如果之前没安装的,也要在anaconda中安装

# 安装 PIL
conda install pillow

(2)下载yolo3源代码
在keras-yolo3的github上下载源代码(https://github.com/qqwweee/keras-yolo3),使用git进行clone或者直接下载成zip压缩文件。

(3)导入PyCharm
打开PyCharm,新建项目,将keras-yolo3的源代码导入到PyCharm中

4、下载预训练模型
YOLO官网上提供了YOLOv3模型训练好的权重文件,把它下载保存到电脑上。下载地址为https://pjreddie.com/media/files/yolov3.weights

5、训练模型
接下来到了关键的步骤:训练模型。在训练模型之前,还有几项准备工作要做。
(1)转换标注数据文件
YOLO采用的标注数据文件,每一行由文件所在路径、标注框的位置(左上角、右下角)、类别ID组成,格式为:image_file_path x_min,y_min,x_max,y_max,class_id
例子如下:
 
这种文件格式跟前面制作好的VOC_2007标注文件的格式不一样,Keras-yolo3里面提供了voc格式转yolo格式的转换脚本 voc_annotation.py
在转换格式之前,先打开voc_annotation.py文件,修改里面的classes的值。例如本案例在voc_2007中标注的熊猫的物体命名为panda,因此voc_annotation.py修改为:

import xml.etree.ElementTree as ET
from os import getcwd

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]

classes = ["panda"]

新建文件夹VOCdevkit/VOC2007,将熊猫的标注数据文件夹Annotations、ImageSets、JPEGImages放到文件夹VOCdevkit/VOC2007里面,然后执行转换脚本,代码如下:

mkdir VOCdevkit
mkdir VOCdevkit/VOC2007
mv Annotations VOCdevkit/VOC2007
mv ImageSets VOCdevkit/VOC2007
mv JPEGImages VOCdevkit/VOC2007

source activate tensorflow
python voc_annotation.py

转换后,将会自动生成yolo格式的文件,包括训练集、测试集、验证集。

(2)创建类别文件
在PyCharm导入的keras-yolo3源代码中,在model_data目录里面新建一个类别文件my_class.txt,将标注物体的类别写到里面,每行一个类别,如下:

(3)转换权重文件
将前面下载的yolo权重文件yolov3.weights转换成适合Keras的模型文件,转换代码如下:

source activate tensorflow
python convert.py -w yolov3.cfg yolov3.weights model_data/yolo_weights.h5

(4)修改训练文件的路径配置
修改train.py里面的相关路径配置,主要有:annotation_path、classes_path、weights_path

其中,train.py里面的batch_size默认是32(第57行),指每次处理时批量处理的数量,数值越大对机器的性能要求越高,因此可根据电脑的实际情况进行调高或调低

(5)训练模型
经过以上的配置后,终于全部都准备好了,执行train.py就可以开始进行训练。

训练后的模型,默认保存路径为logs/000/trained_weights_final.h5,可以根据需要进行修改,位于train.py的第85行,可修改模型保存的名称。

6、使用模型
完成模型的训练之后,调用yolo.py即可使用我们训练好的模型。
首先,修改yolo.py里面的模型路径、类别文件路径,如下:

class YOLO(object):
    _defaults = {
        "model_path": 'logs/000/trained_weights_final.h5',
        "anchors_path": 'model_data/yolo_anchors.txt',
        "classes_path": 'model_data/my_classes.txt',
        "score" : 0.3,
        "iou" : 0.45,
        "model_image_size" : (416, 416),
        "gpu_num" : 1,
    }

通过调用 YOLO类就能使用YOLO模型,为方便测试,在yolo.py最后增加以下代码,只要修改图像路径后,就能使用自己的yolo模型了

if __name__ == '__main__':
    yolo=YOLO()
    path = '/data/work/tensorflow/data/panda_test/1.jpg'
    try:
        image = Image.open(path)
    except:
        print('Open Error! Try again!')
    else:
        r_image, _ = yolo.detect_image(image)
        r_image.show()

    yolo.close_session()

执行后,可爱的熊猫就被乖乖圈出来了,呵呵

通过以上步骤,我们又学习了基于YOLO来训练自己的目标检测模型,这样在应用中可以结合实际需求,使用SSD、YOLO训练自己的数据,并从中选择出效果更好的目标检测模型。

 

关注本人公众号“大数据与人工智能Lab”(BigdataAILab),然后回复“代码”关键字可获取 完整源代码

 

推荐相关阅读

 

展开阅读全文
加载中
点击加入讨论🔥(26) 发布并加入讨论🔥
打赏
26 评论
20 收藏
3
分享
返回顶部
顶部