从零开始带你一步一步使用 YOLOv3 训练自己的数据

公众号:AI有道(redstonewill)

YOLOv3是比较常见和常用的深度学习目标检测(Object Dection)算法。今天给大家介绍一下如何一步一步使用YOLOv3训练自己的数据集。

一、标注数据集

首先我们需要使用 labelimg 工具来标注图片数据集,例如图片是 .jpg 格式的,用矩形框标注图片中的目标位置,得到 .xml 文件。这里 labelimg 的使用方法就不作介绍了,读者自行查阅。

例如 100 张图片对应得到 100 个 .xml 文件:

得到 .jpg 和 .xml 文件之后,我们还需要把数据集整理成 VOC2007 规定的格式,方便我们使用 YOLOv3 进行训练。

方法很简单,建立如下目录层级:

然后将 100 张 .jpg 图片放入到 VOCdevkit/VOC2007/JPEGImages 目录下,将 100 个 .xml 文件放入到 VOCdevkit/VOC2007/Annotations 目录下。

至此,数据集准备完毕!

二、下载 YOLOv3 源码

在我们的 Ubuntu 系统或者服务器上,使用 git 命令直接下载 YOLOv3 的源码工程:

$ git clone https://github.com/pjreddie/darknet

下载完成之后,进入 darknet 目录,会发现如下:

YOLOv3 使用的是开源的神经网络框架 Darknet53,有 CPU 和 GPU 两种模式。默认使用的是 CPU 模式,如果我们使用 GPU 的话,需要修改 darknet 目录下的 Makefile 文件:

$ vim Makefile

修改的几点如下:

修改的地方有 4 处。其中,第 1 处 GPU=1,CUDNN=1,表示使用 GPU 和 CUDNN,若使用 CPU 的话,置为 0 即可;第 2 处将 NVCC 设为本地安装 nvcc 的实际目录;第 3 处将 COMMON+ 设为本地 cuda 的头文件目录;第 4 处将 LDFLAGS+ 设为 cuda 的库目录。

Makeflie 文件修改完成之后,保存退出。

三、准备数据

将我们之前准备好的包含 100 张 .jpg 图片和 100 个 .xml 文件的 VOCdevkit 文件夹拷贝到 darknet 目录下,替换原来的 VOCdevkit 文件夹。

进入到 darknet/VOCdevkit/VOC2007/ 目录下,新建 test.py 脚本,脚本内容如下:

import os
import random

trainval_percent = 0.8
train_percent = 0.8
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
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)

ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/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()

该 test.py 脚本的作用就是对 100 个样本进行训练集和测试集的划分,trainval_percent = 0.8 表示划分 80% 的样本作为训练集,20% 的样本作为测试集。

在当前目录,运行该脚本:

$python test.py

运行完之后,在 /ImageSets/Main/ 目录下会生成 4 个文件:test.txt、train.txt、trainval.txt、val.txt。其中,trainval.txt 存放的是 80 个训练集的样本名,test.txt 存放的是 20 个测试集的样本名。

返回到 drknet 目录下,新建 voc_label.py 脚本,也可以使用下面命令下载:

wget https://pjreddie.com/media/files/voc_label.py

voc_label.py 脚本内容如下:

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

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

classes = ["person", "car", "bicycle", "train"]


def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(year, image_id):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, image_id), 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()

for year, image_set in sets:
    if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
        os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()

该文件需要修改两个地方,分别是 sets 和 classes,classes 根据实际的类别进行修改。

修改完成之后,保存退出。在 darknet 目录下,运行该脚本:

$ python voc_label.py

运行之后,在当前目录下生成 3 个文件:2007_train.txt, 2007_val.txt, 2007_test.txt,存放的是训练集、测试集图片的路径。

运行下面的命令:

$ cat 2007_train.txt 2007_val.txt  > train.txt

这样,train.txt 即为真正的训练集图片路径,2007_test.txt 即为真正的测试集图片路径。

同时,在 VOCdevkit/VOC2007 目录下也会多生成一个 labels 文件夹,labels 文件夹里存放的就是训练集每个图片包含的类别、矩形框四个坐标信息。这时候数据集正式完成。

四、修改配置文件

1. 修改 voc.data 文件

在 drknet 目录下,打开 cfg/voc.data 文件,根据实际情况作如下修改:

其中,class 是类别,train 和 valid 是之前运行 voc_label.py 得到的 train.txt 和 2007_test.txt。names 改为自己的实际路径,backup 为模型存放的路径。

2. 修改 voc.names 和 coco.names

在 darknet 目录下,打开 data/voc.names 和 data/coco.names 文件,修改成自己的类别:

保存退出!

3. 修改 yolov3-voc.cfg

下面需要修改配置文件,打开 cfg/yolov3-voc.cfg,搜索 yolo, 总共会搜出 3 个含有 yolo 的地方。

每个地方都必须要改 2 处:

  • 参数 filters 由下式计算:3*(5+classes),例如本例中 classs=4,则filters=27;

  • 参数 class 改为实际的类别个数;

yolov3-voc.cfg 文件中修改 filters 和 classes 的地方一共有 3 处,注意别遗漏了。

在 yolov3-voc.cfg 文件的开头处,修改训练的一些参数

一般根据具体情况作适当修改即可,注意,训练的时候,Testing 的 batch 和 subdivisions 需要注释掉,learning_rate 是学习率,max_batches 是最大迭代训练次数,可根据训练集大小自行修改。

关于 yolov3-voc.cfg 中更详细的参数解释,可自行查阅。

五、编译

刚才我们修改了 Makefile文件,并修改了各个配置文件,下面就对该工程进行编译,在 darknet 目录下输入下面的命令:

$ make

编译完成之后,生成 darknet 可执行程序。

六、训练

在 darknet 目录下,

首先下载 YOLOv3 的预训练模型:

$ wget https://pjreddie.com/media/files/darknet53.conv.74

然后直接输入下面的命令进行训练:

$ ./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg darknet53.conv.74 -gpus 0,1

其中,-gpus 0,1 表示使用 2 块 gpu 进行训练,编号是 0 和 1。根据你可使用的 gpu 个数进行调整。

当看到上面的打印内容,基本就表示训练开始了,慢慢等待吧~

关于 YOLOv3 详细的训练步骤就介绍到这里,下一篇我将带大家来学习一下使用 YOLOv3 训练好的模型进行实际的图片测试还有如何进行批量测试~


未经允许不得转载:红色石头的个人博客 » 从零开始带你一步一步使用 YOLOv3 训练自己的数据

赞 (2) 打赏

评论 1

  • 昵称 (必填)
  • 邮箱 (必填)
  • 网址
  1. 容添下麻雀虽小五脏俱全回复

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

微信扫一扫打赏