Train YOLO on our own dataset

之前到手 TX1 之后试了一下 YOLO 的 Demo, 感觉很是不错, 帧数勉强达到实时要求, 因此就萌生了使用自己的数据集来训练看看效果的想法.

Dataset

Get The ImageNet Data

为了最大限度地利用资源 (其实是为了偷懒, 但是之后发现给自己挖了个大坑), 我用的是从 ImageNet 上的图片与 Bounding Box 标注. 本次使用了两个类别, 分别是 ballgoal.

  • Downloads内可以可以下到images in the synset以及Bounding Boxes
  • ImageNet 里的图片看起来多, 实际上摊到每个子类上的就1000多张, 能下下来的就500多, 能直接和 Bounding Box 标注匹配的只剩此案100多了TAT. 果然还是需要自己标注, 自力更生.
  • 另外 ImageNet 上的 Bounding Box 信息只有当前类别的. 比如说我下了 goal 的 Bounding Box, 其实某张图片里还有 ball 等 Object, 但是并不会被标出来. 这对于之后的训练有一定影响.
  • 注意如果要从 ImageNet 上下原始图片的话是需要注册账号, 并且通过邮箱认证的 (还不能是 Gmail 这类的可以免费注册的邮箱, 需要机构或者学校邮箱才行).

下好的文件结构如下:

1
2
3
4
5
6
7
.
├── Annotation
│   ├── n03820318
│   └── n04254680
└── images
   ├── n03820318
   └── n04254680

其中Annotation目录下放标注, images目录下放图片.

Convert labels for darknet

ImageNet 上下下来的 Bounding Box 信息是 Pascal VOC 的 xml 格式:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
<annotation>
<folder>n03820318</folder>
<filename>n03820318_101</filename>
<source>
<database>ImageNet database</database>
</source>
<size>
<width>500</width>
<height>333</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>n03820318</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>19</xmin>
<ymin>43</ymin>
<xmax>499</xmax>
<ymax>214</ymax>
</bndbox>
</object>
</annotation>

而 darknet 需要的标注文件是 txt 格式:

1
<object-class> <x> <y> <width> <height>

于是就需要对于 labels 进行转换. 我写了一份 Python 脚本, 将其放在数据集的根目录下执行即可:

1
$ python imagenet_bb_label.py

之后得到目录结构如下:

1
2
3
4
5
6
7
8
9
10
11
12
.
├── Annotation
│   ├── n03820318
│   └── n04254680
├── imagenet_bb_label.py
├── images
│   ├── n03820318
│   └── n04254680
├── labels
│   ├── n03820318
│   └── n04254680
└── train.txt

其中, labels目录保存着转换后的 Bounding Box 信息, train.txt则包含了所有图片文件的绝对路径.

ImageNet 的 xml 文件里 object 的名字是类似n03820318这种格式的, 如果需要转成goal这样的话可以再目录下执行以下命令来批量替换:

1
2
> $ find . -name "*.xml" -print | xargs sed -i 's/<name>n03820318/<name>goal/g'
>

Modify darknet

由于 class 的数量和名字都变了, 因此需要修改下 YOLO 的源代码.

首先是从 clone repository. 可以选择 clone 官方的, 也可以直接下我修改好的. 此处里官方 repo 为例.

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

由于最新的 commit 修改了 label image 的显示方法, 并且改变了源文件里类别的定义, 因此需要先切回之前的 commit:

1
$ git checkout 73f7aacf35ec9b1d0f9de9ddf38af0889f213e99

首先修改Makefile:

1
2
3
GPU=1
CUDNN=1
OPENCV=1

之后是src/yolo.c, 主要是类别名称和数量, 以及train.txtbackup的地址.(backup目录用来存放训练得到的weights)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
// ...
#define NUM_CLASS 2
char *voc_names[] = {"ball", "goal"};
image voc_labels[NUM_CLASS];
void train_yolo(char *cfgfile, char *weightfile)
{
char *train_images = "/home/m/workspace/dataset/train.txt";
char *backup_directory = "/home/m/workspace/backup/";
}
// ...
void test_yolo(char *cfgfile, char *weightfile, char *filename, float thresh)
{
// ...
draw_detections(im, l.side*l.side*l.n, thresh, boxes, probs, voc_names, voc_labels, NUM_CLASS);
// ...
}
void run_yolo(int argc, char **argv)
{
// ...
for(i = 0; i < NUM_CLASS; ++i){
// ...
}
// ...
else if(0==strcmp(argv[2], "demo")) demo(cfg, weights, thresh, cam_index, filename, voc_names, voc_labels, NUM_CLASS, frame_skip);
}

接着是yolo_kernels.cu:

1
2
3
4
5
6
7
8
9
10
// ...
#define NUM_CLASS 2
// ...
void *detect_in_thread(void *ptr)
{
// ...
draw_detections(det, l.side*l.side*l.n, demo_thresh, boxes, probs, voc_names, voc_labels, NUM_CLASS);
// ...
}
// ...

然后是cfg(建议新建一个, 我的配置可作为参考):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# ...
[connected]
# output = Side x Side x (2x5 + class_num)
output= 588
activation=linear
[detection]
# modify the class num
classes=2
coords=4
rescore=1
side=7
num=2
softmax=0
sqrt=1
jitter=.2
object_scale=1
noobject_scale=.5
class_scale=1
coord_scale=5

最后, 如果使用了新的 class 的话, 需要在data/labels里修改make_labels.py并执行来生成新的 label image:

1
2
3
4
5
6
import os
l = ["ball", "goal"]
for word in l:
os.system("convert -fill black -background white -bordercolor white -border 4 -font ubuntu-mono -pointsize 18 label:\"%s\" \"%s.png\""%(word, word))

至此前期准备完成, 可以开始训练了.

Training

首先还需要下载 pre-trained weights.

之后就是慢慢训练之路了:

1
2
3
$ cd darknet
$ make -j8
$ ./darknet yolo train cfg/tiny-yolo.train.cfg darknet.conv.weights

我用上述的 dataset 训练 tiny-YOLO, 从 22:43 一直到 05:12, 总计 6 个小时左右, 最终得到tiny-yolo_final.weights文件.

之后就可以拿来 test 或者 demo 了:

1
$ ./darknet yolo test cfg/tiny-yolo.train.cfg tiny-yolo_final.weights

or

1
$ ./darknet yolo demo cfg/tiny-yolo.train.cfg tiny-yolo_final.weights

Results

yolo-tiny_on_TX1_ball

yolo-tiny_on_TX1_goal

Reference