本文将带你打造一个汽车检测器,使用的算法是PyTorch版本的YOLOV3。本文不会讲解该算法的细节,而是专注于如何去实现自己的汽车检测器,主要包括数据下载,数据清洗,数据集制作以及训练和检测(图片检测,视频检测)等过程。即使你不知道YOLOV3是什么也没关系,这对于阅读本文几乎毫无影响。
Alt text

数据集下载

为了检测汽车,首先需要获取含有汽车的标注好的数据集,这里我选择了KITTI。由于官网下载太慢,推荐使用下面的百度云链接进行下载。

下载KITTI数据集:
https://pan.baidu.com/s/1t3TXXkqVR4NGqZwQiGEIzg
提取码:cw35

源码来自Github开源项目https://github.com/eriklindernoren/PyTorch-YOLOv3,为了适应本文的内容,我将其做了些许更改,一并上传到了百度云,在公众号后台回复"detect"直达下载地址。

数据清洗

在主目录PyTorch-YOLOv3-master下新建文件夹datasets,在datasets中新建四个文件夹:train_image,train_label,proceed_train_image_path,proceed_train_label_path

将下载好的图片数据(.png)全部移动到train_image,标签数据(.txt)全部移动到train_label.

Alt text

本部分的数据清洗主要针对以下问题:

  • 由于不可知的原因,下载下来的数据中,有些图片是损坏的,不能够被加入训练集

  • 原始kitti数据集的标注标签中,除了模型必须的信息(1个目标所属类别+4个位置坐标)外,还包含其余信息,我们需要过滤掉这些信息;

  • 同时,我们只需要检测Car,因此需要将其余类别(非Car)的信息也过滤掉。

  • 在kitti数据集中,坐标格式为:<left top right bottom>
    Alt text

而本次使用的YOLOV3模型要求的格式为:<x_center y_center width height>
Alt text

因此需要做一个转换,转换公式为:
$$width=right-left$$
$$height=bottom-top$$
$$x_{center}=left+width/2$$
$$y_{center}=top+height/2$$

  • 根据模型的要求,坐标需要做归一化处理:
    $$width <— width/w$$
    $$height <— height/h$$
    $$x_{center} <— x_{center}/w$$
    $$y_{center} <— y_{center}/h$$

其中,$w$和$h$分别代表训练集图片的宽度和高度。

实现上述要求的代码如下:

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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import os
import cv2
import shutil

#原数据存放路径
train_img_path='train_image'
train_label_path='train_label'
#过滤后的数据存放路径
proceed_train_image_path='proceed_train_image_path'
proceed_train_label_path='proceed_train_label_path'
#图像尺寸
w=1242
h=375

for img_name in os.listdir(train_img_path):
current_img=cv2.imread(os.path.join(train_img_path,img_name))
if current_img is not None:
print(img_name)
label_name=img_name.split('.')[0]+'.txt'
label_path=os.path.join(train_label_path,label_name)

flag=0#当前图片是否包含Car
with open(label_path,'r') as f:
all_obj=f.readlines()
#对于当前图片中每一个目标
for obj in all_obj:
info=obj.split(' ')
name=info[0]#类别
left,top,right,bottom=float(info[4]),float(info[5]),float(info[6]),float(info[7])#四个坐标
#将四个坐标转换为x_center,y_center,width,height的格式
width=right-left
height=bottom-top
x_center=left+width/2
y_center=top+height/2
#将其归一化
width=width/w
height=height/h
x_center=x_center/w
y_center=y_center/h

#发现car,开始写入
if name == 'Car':
flag=1
with open(os.path.join(proceed_train_label_path,label_name),'a+') as nf:
nf.write('0 '+str(x_center)+' '+str(y_center)+' '+str(width)+' '+str(height)+'\n')

if flag==1:
pass
shutil.copy(os.path.join(train_img_path,img_name),os.path.join(proceed_train_image_path,img_name))

else:
print('{}图片已损坏'.format(img_name))

datasets文件夹下打开jupyter notebook,新建文件,写入上述代码并运行,就完成了数据清洗。

清洗后的图片数据(.png)存放在proceed_train_image_path,与之对应的标签数据(.txt)存放在proceed_train_label_path

其中,图片数据就是一张张的图片,而标签数据是一个个的文本文档,并且每一张图片对应一个这样的文本文档:
Alt text

文本文档的内容是一行一行的数据,每一行都代表该文本文档对应的图片中被标注的一个目标(在这里的目标就是车, Car),图片中有几辆车被标注出来,该文本文档就有多少行。

每一行包含5个信息,第一个是目标所属类别,这里由于只检测车(Car),也就是只有1类,因此都是0(根据模型的要求,0代表Car,我们只需在data/custom/classes.names中的第一行写入Car即可),后面四个分别代表了处理后(坐标转换+归一化)的边界框的坐标:
Alt text

制作数据集

将上一步得到的图片数据(.png)和标签数据(.txt)分别移动到data/custom/imagesdata/custom/labels,若没有这两个文件夹则需先手动创建。

现在开始制作数据集。主要是制作训练集和验证集对应的.txt文件:

在主目录PyTorch-YOLOv3-master下,打开jupyter notebook,新建文件,写入以下代码:

1
2
3
4
5
6
7
8
9
10
import os
img_path='data/custom/images'
#制作.txt
for i,img_name in enumerate(os.listdir(img_path)):
if i<100:
with open(os.path.join('data/custom','valid.txt'),'a+') as f1:
f1.write(os.path.join(img_path,img_name+'\n'))
else:
with open(os.path.join('data/custom','train.txt'),'a+') as f2:
f2.write(os.path.join(img_path,img_name+'\n'))

从上述代码可以看出,选取了前100张作为验证集,其余作为训练集。

现在,data/custom路径下的所有文件(夹)结构如下:
Alt text

开始训练

在训练之前,先将作者提供的预训练好的darknet53权重下载下来,具体地,在weights文件夹中,执行如下命令:

1
sh download_weights.sh

它还会额外下载另外两个权重:yolov3.weightsyolov3-tiny.weights,这些其实是用不到的,可以忽略。

然后进入config文件夹下,执行如下命令:

1
sh create_model.sh 1

这是在制作yolov3网络的配置文件,其中的1代表总共有几类,由于这里只检测车(Car),所以只有1类。

完成以上操作后,如果显存足够大,就可以开始训练了,在主目录PyTorch-YOLOv3-master下打开终端,执行以下命令:

1
python train.py --model config/yolov3-custom.cfg --data config/custom.data --pretrained_weights weights/darknet53.conv.74

如果显存较小,可能会报CUDA内存溢出的错误,可以手动调小batch_size。我是在GTX1070上进行训练的,就报了这个错误,想调一下batch_size,却发现它不是直接给出的一个超参数,而是计算得到的,打印了一下,发现当前的batch_size为16:
Alt text

我采取了一个简单粗暴的方法,就是直接手动将mini_batch_size设置为4:
Alt text

现在,再次执行上面的命令,就开始了了漫长的训练过程
Alt text

开始检测

训练完成后,就可以使用该模型对自己的数据进行检测了。

将待检测的图片数据放在data/samples,注意先清空这个文件夹下的其他图片:
Alt text

然后,打开detect.py,修改其中run函数的部分代码:
Alt text

weights中的141代表这是第141个epoch得到的权重,你可以自行选择效果较好的epoch对应的权重文件,只需将141替换为你的epoch即可。

现在,在主目录PyTorch-YOLOv3-master下,打开终端,输入如下命令:

1
python detect.py --images data/samples/

检测过程输出如下:
Alt text

检测完成后,在output文件夹中,就能看到检测后的结果,每一个被检测出来的Car都被框起来了:
Alt text

如果你已经完成了训练,但是在检测过程中没有得到上述输出,在output文件夹中的图片中的Car也没有被框起来,可以打开detect.py,将conf_thres调小一些,比如我调成了0.1:
Alt text

之后保存该文件,再次执行检测的命令就可以了。

视频检测

将视频切分成帧,每一帧都是一张图片,将全部的帧送入检测器进行检测,再将检测后的图片合成回视频,就完成了对视频的目标检测。

在主目录PyTorch-YOLOv3-master下,新建文件夹process_videoimage,将待检测的视频放到process_video文件夹中,并在process_video中新建脚本,写入如下的视频切分代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
import cv2
vc=cv2.VideoCapture("segment-1.avi")
c=1
if vc.isOpened():
rval,frame=vc.read()
else:
rval=False
while rval:
rval,frame=vc.read()
cv2.imwrite('image/'+str(c)+'.jpg',frame)
c=c+1
cv2.waitKey(1)
vc.release()

运行上述代码,就完成了视频的切分,切分后的图片存放在新建的image文件夹中。

执行如下检测命令:

1
python detect.py --images process_video/image

检测后的图片被保存在output文件夹中,同样注意检测前先清空该文件夹,便于接下来的视频合成。

检测出来的部分结果展示如下:
Alt text

现在,在主目录PyTorch-YOLOv3-master下,新建脚本,键入如下代码,即可将图片合成为视频:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import cv2
import os
img=cv2.imread('output/10.png')
size=img.shape[:2][::-1]
fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
videoWriter = cv2.VideoWriter('out.avi', fourcc, 30, size)
img_dir='output'
for img_name in sorted(os.listdir(img_dir),key=lambda x: int(x.split('.')[0])):
print(img_name)
img = os.path.join('output', img_name)
frame = cv2.imread(img)
videoWriter.write(frame)
videoWriter.release()

输出的视频文件out.avi保存在主目录PyTorch-YOLOv3-master下。