如何训练YOLOv2来识别自己的对象

概要

我的系统环境

OS: Ubuntu 16.04
Python Version: 2.7
GPU: 1gb

在本篇文章中我们想要识别的对象是NFPA 704 ‘fire diamond’.

本文的图片集可以在这里下载

2安装Darknet

ubuntu:~$ git clone https://github.com/AlexeyAB/darknet.git
ubuntu:~$ cd darknet

（如果你需要GPU加速请按照以下步骤操作.GPU加速对训练很重要,如果不想训练则不需要）

ubuntu:~$ vi Makefile

(修改GPU 0 为 1并保存.如果你安装了OpenCV,设置OpenCV 0 为 1,否则不需要)¹

ubuntu:~$ make

如果有报错请尝试解决,如果编译一切正常,试着运行它!

./darknet

你应该会得到如下的输出

usage: ./darknet <function>

3测试Yolo

在子文件夹cfg/中已经含有了YOLO的config文件.你现在要下载(258MB)已经训练好的weight文件.或者直接输入下面的代码下载:

wget https://pjreddie.com/media/files/yolo.weights

然后尝试使用已训练好的weight识别!

./darknet detect cfg/yolo.cfg yolo.weights data/dog.jpg

你会得到一些像这样的输出:

layer     filters    size              input                output
    0 conv     32  3 x 3 / 1   416 x 416 x   3   ->   416 x 416 x  32
    1 max          2 x 2 / 2   416 x 416 x  32   ->   208 x 208 x  32
    .......
   29 conv    425  1 x 1 / 1    13 x  13 x1024   ->    13 x  13 x 425
   30 detection
Loading weights from yolo.weights...Done!
data/dog.jpg: Predicted in 0.016287 seconds.
car: 54%
bicycle: 51%
dog: 56%

如果你想要获得更多关于darknet的信息请点击这里.

4给图片添加释文

这一步主要是为图片创建文本文件（只需要理解下面的步骤）

1.在与.jpg图片相同的路径下为每个图片创建与图片同名的.txt文件.在创建.txt文件的同时,将以下信息写入.txt文件:要识别对象的数量和对象在图片中的坐标，每个对象占一行：<object-class> <x> <y> <width> <height>

在上述代码中：

1. 是对象的种类，值为0到(classes-1).
2. 是长、宽相对于图片的小数值，在0.0到1.0的范围之内。
3. 例如<x> = x在图片中的绝对坐标 / 图像的宽 ，<height> = 对象所在区域的高度 / 图像的高度
4. 注：<x> <y>是对象所在长方形区域的中心（不是左上角）

例如，为img1.jpg创建的img1.txt文件需要包含如下内容：

1 0.716797 0.395833 0.216406 0.147222
0 0.687109 0.379167 0.255469 0.158333
1 0.420312 0.395833 0.140625 0.166667

例如：

我的图片尺寸为: 360 * 480，其中包含一个对象：狗
image_width = 360
image_height = 480
absolute_x = 30 (狗在图像中的x坐标)
absolute_y = 40 (狗在图像中的y坐标)
absolute_height = 200 (狗在图像中所占长方形区域的高)
absolute_width = 200 (狗在图像中所占长方形区域的宽)

那么最后.txt文件需要有5个参数，并且五个参数都在0到1之间

<class_number>  (<absolute_x> / <image_width>)  (<absolute_y> / <image_height>)  (<absolute_width> / <image_width>)  (<absolute_height> / <image_height>)
那么：
0 (30/360) (40/480) (200/360) (200/480)
最后文件中有的是如下五个小数：
0 0.0833 0.0833 0.556 0.417

5开始训练

注：请先待训练这个只含有一个对象的数据集成功后，再开始训练你自己的(含有多个对象的)数据集
我有了一个新的标注工具，请在这里下载新的工具来准备数据。如果你使用新的标注工具，在本博客的例子中请不要运行main.py和convert.py，直接运行process.py。因为两个python文件进程采用单个python脚本.需要完全的介绍请看新标注工具的介绍视频。

在这里下载已经标注好的图片据集。
在这里下载未标注的图片据集。

如果你下载了未标注的数据集，你可以用以下命令生成标注文件：
从GitHub上下载代码
把你所有的图片放进Images/001,然后运行main.py
然后你就可以使用这个工具对图片进行标注。
图片标注工具的使用介绍视频在这里。
但因为需要梯子，所以将以截图加叙述的方式进行说明。

1. 清空Images/001文件夹下的文件，并把下载的图片集放进去
2. 删除Labels文件夹下的所有内容，并在Labels文件夹下新建两个文件夹，分别名为001，和名为output
3. 然后打开终端，进入main.py所在目录，运行main.py
4. 打开软件后设置Image Dir为001并点击load，如下
5. 用鼠标选中需要识别的区域
6. 点击next直到把所有图片标注完成

使用convert.py生成最终的YOLO使用的txt文件，运行convert.py，在这里可以看到视频介绍如何使用convert.py。不过需要梯子。同样的，本文以文字进行讲解。

1. 打开终端，运行convert.py，会自动的将Labels/001下面的文件转换后写入到Labels/output文件夹中。

通过比较转换后和转换前的文件可以发现，BBox 标签格式为

   [类型号]
    [左边的x坐标] [上面的y坐标] [右边的x坐标] [下方的y坐标]

而 YOLOv2的格式为

[类型号] [中心x坐标除以整个图片的宽] [中心的y坐标除以整个图片的高] [对象所在区域的宽除以整个图片的宽] [对象所在区域的高除以整个图像的高]

通过修改convert.py第31和32行可以改变convert.py读取文件和输出文件的位置。

太好了！现在每个图片都有了自己的.txt文件了，这些文件会告诉YOLO哪里是我们需要识别的目标。
到现在为止，数据集已经完全准备好了。将刚才output中的文件放入图片集所在的文件夹中，文件夹将会看起来像这样：

下一步，我们需要test.txt文件和train.txt文件来告诉YOLO哪些图片是用来训练的训练集，哪些图片是用来测试的测试集。
我写了一个小的python脚本process.py来生成上面两个文件，这个python脚本会在它运行的目录下寻找图片(将它放在数据集的文件夹下，和图片与.txt文件放在一起)。你也可以修改第五行来指向其他的路径。
process.py代码如下²

import glob, os# Current directory
current_dir = os.path.dirname(os.path.abspath(__file__))print(current_dir)
#current_dir = '<Your Dataset Path>'# Percentage of images to be used for the test set
percentage_test = 10;# Create and/or truncate train.txt and test.txt
file_train = open('train.txt', 'w')  
file_test = open('test.txt', 'w')# Populate train.txt and test.txt
counter = 1  
index_test = round(100 / percentage_test)  
for pathAndFilename in glob.iglob(os.path.join(current_dir, "*.jpg")):  
    title, ext = os.path.splitext(os.path.basename(pathAndFilename))if counter == index_test:
        counter = 1
        file_test.write(current_dir + "/" + title + '.jpg' + "\n")
    else:
        file_train.write(current_dir + "/" + title + '.jpg' + "\n")
        counter = counter + 1

运行上面的代码就会得到如下的两个文件：

6准备YOLOv2的配置文件

YOLOv2需要明确的文件来指定如何训练和训练什么。我们将会创建三个文件，我是用的是1GB内存的GPU所以我是用tiny-yolo.cfg文件。三个文件分别如下：

cfg/obj.data
cfg/obj.names
cfg/tiny-yolo.cfg

首先我们创建obj.data文件，这个文件会说明我们的目标有多少种，train.txt文件在哪里，test.txt文件在哪里，.names文件在哪呢，以及指定训练好的网络将保存在哪里。
我们这次的obj.text文件内容如下

classes= 1  
train  = train.txt  
valid  = test.txt  
names = obj.names  
backup = backup/

我们这次的obj.names文件如下，每一个种类的对象名称都应该换行。并且总的行数应该和我们之前给图片标注标签时的对象种类数相匹配。

NFPA

最后我们需要准备（我知道，你的GPU早已饥渴难耐了！）的是.cfg文件，我直接使用cfg/tiny-yolo.cfg文件，并做少量的改动。改动如下

1. 第2行:设置batch=24，这意味着每一步的训练会使用64张图³
2. 第3行，设置subdivisions=8，如果你的强劲的GPU有很多内存，你可以减小这个值，或者增大batch
3. 第120行，设置classes=1，这个值为我们有多少中东西想要识别，本文为1种
4. 第114行，设置filters=(classes + 5)*5，本文为filters=30

YOLOv2需要一个卷积权重设置文件来开始训练。为了让事情变得简单，Joseph提供了一个已经用Imagenet训练过的文件。这个文件给我们提供了一个极好的初始化权重值。在这里下载。完成下载后我们即可进行下一步。

7训练

输入以下命令，然后静静地观看你的GPU做她最喜欢做的事（将你生成的train.txt和test.txt文件复制到yolo_darknet的目录下）：

manivannan@manivannan-whirldatascience:~/YoloExample/darknet$ ./darknet detector train cfg/obj.data cfg/yolo-obj.cfg darknet19_448.conv.23

备注：当迭代达到100时，yolo会自动保存一次训练处的文件，并且当lose下降到小于0.06时，yolo会自动停止训练。

8结果

我们应该会在中得到一个.weights文件，那就是我们训练的模型。让我们来用一些图测试一下识别NFPA 704 ‘fire diamond’ 的效果。以下命令会使YOLOv2识别我们指定的图片。如何你想要我训练好的模型，可以在这里下载。

manivannan@manivannan-whirldatascience:~/YoloExample/darknet$ ./darknet detector test cfg/obj.data cfg/yolo-obj.cfg yolo-obj1000.weights data/manivannan.jpg

我收集了一些有趣的图片来证明我们的成果和识别设置。这些图片没有在训练中使用过。正如你所见，下图中的每个NFPA ‘fire diamond’ 都被识别出来了。完美！

本文翻译自https://medium.com/@manivannan_data/how-to-train-yolov2-to-detect-custom-objects-9010df784f36。

---

1. 译者注：如果想用opencv加载图片，更改opencv为1，我通常opencv为0. ↩︎

2. 译者注：在此代码中index_test控制每几张图片拿出一张作为测试集 ↩︎

3. 这里可能是作者笔误，应该是24张图？ ↩︎