YOLOv3+Tensorflow训练自己的数据

最近在做一个用YOLOv3进行目标检测的项目，根据网上各位大神的博客，结合自己做项目的过程做了一个笔记，方便自己以后回顾实验过程，也给大家做个参考。
实验用的程序主代码来自于github上一位大神程序猿，这里给出代码的地址
YOLOv3-tensorflow大神源码

一、制作数据集

根据代码来看，训练和测试使用的数据格式需要按照PascalVOC的数据格式来，所以首先我们需要进行数据集的制作。

1 获取数据

根据实验需要，我将录制的视频文件分为训练视频和测试视频，然后按帧进行截取，获得训练和测试用的图片，因为实验的保密性，在此不能明具体的内容。将视频按帧截取为图片并保存使用的是下面这段代码：

import cv2
import glob
import os
from datetime import datetime

def video_to_frames(path):
    """
    输入：path(视频文件的路径)
    """
    # VideoCapture视频读取类
    videoCapture = cv2.VideoCapture()
    videoCapture.open(path)
    for i in range(int(frames)):
        ret, frame = videoCapture.read()
        cv2.imwrite("E:\dataset\images\train%d.jpg" % (i), frame)
    return

if __name__ == '__main__':
    video_to_frames("E:\dataset\video\train.mp4")
    print("SUCCEED !!!")

2 标记图片

根据PascalVOC数据集的需要，使用Labelimg工具对图片进行标注，标注后会生成XML文件，如下图所示：

2 按照PascalVOC数据集的格式整理自己的数据

这次实验我分别建立了VOCTrainval和VOCTest两个数据文件，分别用于训练和测试，大家也可以不分开，后面进行训练和测试数据划分就行了，两个文件夹都按照PascalVOC的格式建立。PascalVOC数据集包含了5个部分，在实验中我们只需要用到一下三个文件夹：
1） Annatations文件夹
文件夹存放的是xml格式的标签文件，每个xml文件都对应于JPEGImages文件夹的一张图片。
2）JPEGImages文件夹
文件夹里包含了训练图片或测试图片。
3)ImageSets文件夹
该文件夹里原有三个子文件夹，但实验中我们仅需要使用Main文件夹里面的信息，存放的是图像物体识别的数据，有train.txt, val.txt ,trainval.txt.这三个文件（VOCTrainval文件夹下）或者test.txt 文件（VOCTest）。这几个文件我们后面会生成。
按照要求，将自己的图片放入JPEGImages文件夹，将标注信息xml文件放入Annatations文件夹：

3 划分训练集和测试集

训练时要有测试集和训练集，如果在制作数据集的时候没有像我一样进行区分，那么在这里就需要使用代码将数据进行划分，放在ImageSets\Main文件夹下。代码如下，至于训练验证集和测试集的划分比例，以及训练集和验证集的划分比例，根据自己的数据情况决定。将下面的代码放入split.py中：

import os  
import random   
  
xmlfilepath=r'E:/tensorflow-yolov3-master/VOCData/VOCTrainVal/Annotations/'  #xml文件的路径
saveBasePath=r'E:/tensorflow-yolov3-master/VOCData/VOCTrainVal/ImageSets/' #生成的txt文件的保存路径
  
trainval_percent=0.9 
train_percent=0.8  
total_xml = os.listdir(xmlfilepath)  
num=len(total_xml)    
list=range(num)    
tv=int(num*trainval_percent)    
tr=int(tv*train_percent)    
trainval= random.sample(list,tv)    
train=random.sample(trainval,tr)    
  
print("train and val size",tv)  
print("traub suze",tr)  
ftrainval = open(os.path.join(saveBasePath,'Main/trainval.txt'), 'w')    
ftest = open(os.path.join(saveBasePath,'Main/test.txt'), 'w')    
ftrain = open(os.path.join(saveBasePath,'Main/train.txt'), 'w')    
fval = open(os.path.join(saveBasePath,'Main/val.txt'), 'w')    
  
for i  in list:    
    name=total_xml[i][:-4]+'\n'    
    if i in trainval:    
        ftrainval.write(name)    
        if i in train:    
            ftrain.write(name)    
        else:    
            fval.write(name)    
    else:    
        ftest.write(name)    
    
ftrainval.close()    
ftrain.close()    
fval.close()    
ftest .close() 

直接运行split.py，或者对路径和划分比例稍做修改后分别对VOCTrainval和VOCTest文件夹中的数据运行即可，得到ImageSets\Main\文件下的几个txt文件。到此数据集制作的第一个阶段完成！

二、在data/dataset文件夹下生成数据的描述文件

根据作者在github上的声明，该程序在训练时还需要两个对图片信息的描述文件voc_train.txt和voc_test.txt，格式如下

作者给出了生成这两个文件的代码scripts\voc_annotation.py。需要做部分的修改：

1）将classes修改为自己的类别

classes = ['类别1','类别2','类别3','类别4']

2. 将相关的路径修改为自己的路径

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 将相关路径中default的路径修改为自己的路径
    parser.add_argument("--data_path", default="E:/tensorflow-yolov3-master/VOCData/")
    # 将default的路径更改为自己的voc_train.txt存放的位置
    parser.add_argument("--train_annotation", default="../data/dataset/voc_train.txt")
     # 将default的路径更改为自己的voc_test.txt存放的位置
    parser.add_argument("--test_annotation",  default="../data/dataset/voc_test.txt")
    flags = parser.parse_args()

    if os.path.exists(flags.train_annotation):os.remove(flags.train_annotation)
    if os.path.exists(flags.test_annotation):os.remove(flags.test_annotation)

    # 更改训练集和测试集的相对路径，由于本实验的训练数据只有一个路径，所以注释掉了num2
    num1 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTrainVal'), 'trainval', flags.train_annotation, False)
    #num2 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTrainVal'), 'trainval', flags.train_annotation, False)
    num3 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTest'),  'test', flags.test_annotation, False)
    print('=> The number of image for train is: %d\tThe number of image for test is:%d' %(num1, num3))

做了以上更改后直接运行该文件即可。这时候data\dataset文件加下将多出两个文件：

查看其中一个文件的内容如下：

三、修改其他文件

1 修改voc_names

打开data/classes/文件夹中的voc_names文件，将其中的类别名称改为自己的类别。

2 修改配置文件

打开core\文件夹下的config.py文件

根据自己的文件路径，进行适当修改，当然，若按照我上述步骤来的，路径应该是没有什么需要修改的地方的，那些参数也自己按照训练的过程慢慢调就行了。

四、训练

接下来就是训练啦，直接运行train.py就行了，一般来说是不会有什么问题的。

五、测试

训练完成之后，终于到了最激动人心的时刻了，进行模型的测试！

1 修改config文件

修改__C.TEST.ANNOT_PATH为自己voc_test.txt文件的路径
修改__C.TEST.WEIGHT_FILE为自己需要测试的模型的路径
__C.TEST.INPUT_SIZE根据需要自行选择，也可以不做修改

# TEST options
__C.TEST                        = edict()

__C.TEST.ANNOT_PATH             = "./data/dataset/voc_test.txt"  #修改为自己voc_test.txt文件的路径
__C.TEST.BATCH_SIZE             = 6
__C.TEST.INPUT_SIZE             = 544  #自行决定是否修改，不修改也没有影响
__C.TEST.DATA_AUG               = False
__C.TEST.WRITE_IMAGE            = True
__C.TEST.WRITE_IMAGE_PATH       = "./data/detection/"
__C.TEST.WRITE_IMAGE_SHOW_LABEL = True
__C.TEST.WEIGHT_FILE            = "./checkpoint/yolov3.ckpt" # 修改为自己模型的路径
__C.TEST.SHOW_LABEL             = True
__C.TEST.SCORE_THRESHOLD        = 0.3
__C.TEST.IOU_THRESHOLD          = 0.45

2 执行evaluate.py文件

这里没什么需要修改的，我后来是有增加一个保存检测后图片的代码，因为作者给的图片测试代码效果太差了，这个后面再讲。

3 修改mAP\extra\中的class_list.txt文件

打开class_list.txt文件，将内容修改为自己的类别。

4 计算mAP和各类AP

1）若类别名称中有空格，需要先执行mAP\extra\文件夹中的remove_space.py或者rename_class.py，将空格转换成短横‘-’。若类别名中没有空格则忽略这一步。

2）执行mAP\main.py，得出各类AP以及mAP等，其结果保存在mAP\results\目录下。

六、图片测试和视频测试

作者给出的测试代码需要将模型另外保存为.pb的形式，所以需要先对模型进行转换。

1 生成YOLOv3.pb文件

修改freeze_graph.py
1)将pb_file修改为自己希望.pb模型保存的路径

pb_file = "./model_pb/yolov3.pb"

2）将ckpt_file修改为自己的.ckpt模型的路径

ckpt_file = "./checkpoint/yolov3_test_loss=7.8478.ckpt-62"

然后直接运行freeze_graph.py文件。

2 图片测试

修改image_demo.py
将pb_file修改为自己模型的路径
将video_path修改为自己图片的路径

pb_file         = "./model_pb/yolov3.pb"
video_path      = "./docs/images/test1.jpg"

运行image_demo.py
不知道为啥，效果很差，但是在上一步的测试中mAP已经达到了87%，效果不应该会有这么差的，所以我就直接在evaluate.py中增加了几行代码，将evaluate.py中生成的检测后的图片全部保存了下来，果然，检测的效果是很好的，增加的几行代码如下：

if self.write_image:
                    image = utils.draw_bbox(image, bboxes_pr, show_label=self.show_label)
                    cv2.imwrite(self.write_image_path+image_name, image)
                    # 下面三行为新增代码
                    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    image = Image.fromarray(image)
                    image.save(output_path+image_name)

3 视频测试

跟上一步一样修改video_demo.py
将pb_file修改为自己模型的路径
将video_path修改为自己图片的路径

pb_file         = "./model_pb/yolov3.pb"
video_path      = "./docs/images/test.mp4"

跟图片测试一眼，虽然能够成功运行，但是检测的效果很差，另外增加了一些程序运行的信息，发现检测速度也不是很好。后来修改代码，使用帧交替双线程的方法进行测试，发现不仅速度快了，检测效果也变好了，具体是什么原因导致的，目前我也不清楚。帧交替双线程的视频检测代码将在下一篇博客中进行叙述。

本人目标检测初学者，博客中或许有一些不尽完善的地方，欢迎大家提出来一起交流讨论。