使用Python在自定义数据集上训练YOLO进行目标检测
使用Python在自定义数据集上训练YOLO进行目标检测
介绍
最近开始在计算机视觉领域工作。在这些早期日子里，我们正在研究各种目标检测算法的工作原理。其中最知名的算法包括R-CNN、Fast R-CNN、Faster R-CNN和当然是YOLO。
在本文中，重点介绍最后提到的算法。YOLO是目标检测领域的最新技术，有无数的用例可以使用YOLO。然而，今天不想告诉你YOLO的工作原理和架构，而是想简单地向你展示如何启动这个算法并进行预测。此外，我们还将看到如何在自定义数据集上训练它，以便你可以将其适应你的数据。
Darknet
我们认为没有比你可以在他们的网站链接中找到的定义更好地描述Darknet了。
Darknet是一个用C和CUDA编写的开源神经网络框架。它快速、易于安装，并支持CPU和GPU计算。你可以在GitHub上找到源代码，或者你可以在这里了解更多关于Darknet能做什么的信息。
所以我们要做的就是学习如何使用这个开源项目。
你可以在GitHub上找到darknet的代码。看一看，因为我们将使用它来在自定义数据集上训练YOLO。
克隆Darknet
我们将在本文中向你展示的代码是在Colab上运行的，因为我没有GPU…当然，你也可以在你的笔记本上重复这个代码。偶尔会更改路径。
所以首先我们要克隆darknet的GitHub存储库。Colab允许我们使用%%bash命令来编写bash命令。
%%bash
git clone https://github.com/AlexeyAB/darknet
一旦你克隆了存储库，你将在你的工作目录中看到许多文件，放松，它看起来比实际复杂得多。现在我们需要重新配置makefile。不知道什么是makefile吗？简化来说，它是一个使编译代码变得容易的文件。
如果你曾经在C中编写过代码，你知道实践是在写完一个文件file.c之后，使用像g++等命令来编译它… 在大型项目中，这个编译命令可能会非常长，因为它必须考虑到依赖关系等等。
因此，在每次编译时都去重新编写g++等命令将会非常费力… 那么我们要做的是创建一个makefile，它已经包含了这个命令，并且我们所需要做的就是运行makefile来编译代码。makefile通常包含用户可以根据需要设置的配置变量。
话虽如此，我们要做的是设置Darknet makefile中的一些变量。确保你有一个可用的GPU，并运行以下单元格。
%%bash
cd darknet
sed -i 's/OPENCV=0/OPENCV=1/' Makefile
# In case you dont have a GPU, make sure to comment out the
# below 3 lines
sed -i 's/GPU=0/GPU=1/' Makefile
sed -i 's/CUDNN=0/CUDNN=1/' Makefile
sed -i 's/CUDNN_HALF=0/CUDNN_HALF=1/' Makefile
在这个单元格的命令sed -i，例如在第一行中，允许你将OPENCV变量从0更改为1。
我们在上一个单元格中设置的配置允许我们在GPU上启动YOLO，而不是在CPU上。现在我们将使用make命令来启动makefile。
%%bash
#compile darkent source code
cd darknet
现在我们安装一个库，它将用于在YOLO检测到的对象周围绘制边界框。
%%capture
!pip install -q torch_snippets
下载数据集
我们将使用一个包含卡车和公共汽车图像的目标检测数据集。Kaggle上有许多目标检测数据集，你可以从那里下载一个。
如果你不知道如何在Colab中直接从Kaggle下载数据集，你可以去阅读一些我以前的文章。
所以下载并解压数据集。
!wget - quiet link_to_dataset
!tar -xf open-images-bus-trucks.tar.xz
!rm open-images-bus-trucks.tar.xz
下载的数据集的结构如下图所示。
下载YOLO
显然，你不必从头开始训练YOLO，而是可以直接从互联网上下载权重。要从YOLO4下载权重，请使用以下命令。
!wget - quiet https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights
要查看一切是否正常运行，请运行以下命令。
%%bash
#I had to use the flag -dont_show cause wasnt working. Try to run wiithout it
cd darknet
./darknet detector test cfg/coco.data cfg/yolov4.cfg ../yolov4.weights data/person.jpg -dont_show
在此命令中，我们刚刚运行了：
- 我们指定了YOLO4的配置：cfg/yolov4.cfg
- 我们指定要使用的权重：../yolov4.weights
- 我们将在coco数据集上进行预测，因为你克隆了存储库：cfg/coco.data
- 我们对以下图像进行预测：data/person.jpg
准备你的数据集
YOLO期望正确设置某些文件和文件夹，以便在自定义数据集上进行训练。
首先，你需要打开darknet/data/obj.names路径下的文件，其中你将写入你的标签。
在Colab中，我们可以使用魔术命令直接在一个单元格中写入文件。魔术命令下的所有内容都将被复制到指定的文件中。
%%writefile darknet/data/obj.names
bus
truck
现在我们需要修改另一个文件，告诉YOLO期望多少类别，以及在哪里找到训练和验证的路径，以及在哪里找到带有标签名称的文件。我们可以使用magic命令和以下几行来简单地完成这个操作。
%%writefile darknet/data/obj.data
classes = 2
train = darknet/data/train.txt
valid = darknet/data/val.txt
names = darknet/data/obj.names
backup = backup/
为了理解它，你的train txt文件应该像下图所示（验证类似）。
其中每一行指示在哪里找到训练图像。
尽管我们指定的文件仍然是空的。所以我们将这些数据从我们下载的数据集文件夹复制到Darknet默认文件夹中。
!mkdir -p darknet/data/obj
!cp -r open-images-bus-trucks/images/* darknet/data/obj/
!cp -r open-images-bus-trucks/yolo_labels/all/{train,val}.txt darknet/data/
!cp -r open-images-bus-trucks/yolo_labels/all/labels/*.txt darknet/data/obj/
#add prefix 'darkent/' in front of each row in darkent/data/train.txt
!sed -i -e 's/^/darknet\//' darknet/data/train.txt
!sed -i -e 's/^/darknet\//' darknet/data/val.txt
与之前一样，我们将下载YOLO的权重。这次我们下载yolov4-tiny，它比之前的更快。然后我们将权重复制到Darknet内的适当文件夹中。
!wget - quiet https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v4_pre/yolov4-tiny.conv.29
!cp yolov4-tiny.conv.29 darknet/build/darknet/x64/
现在让我们重新命名负责配置yolov4 tiny架构的配置文件。然后，我们将编辑一些参数，设置批次数、类别数和其他参数。
%%bash
cd darknet
# create a copy of existing configuration and modify it in place
cp cfg/yolov4-tiny-custom.cfg cfg/yolov4-tiny-bus-trucks.cfg
# max_batches to 4000 (since the dataset is small enough)
sed -i 's/max_batches = 500200/max_batches=4000/' cfg/yolov4-tiny-bus-trucks.cfg
# number of sub-batches per batch
sed -i 's/subdivisions=1/subdivisions=16/' cfg/yolov4-tiny-bus-trucks.cfg
# number of batches after which learning rate is decayed
sed -i 's/steps=400000,450000/steps=3200,3600/' cfg/yolov4-tiny-bus-trucks.cfg
# number of classes is 2 as opposed to 80
# (which is the number of COCO classes)
sed -i 's/classes=80/classes=2/g' cfg/yolov4-tiny-bus-trucks.cfg
# in the classification and regression heads,
# change number of output convolution filters
# from 255 -> 21 and 57 -> 33, since we have fewer classes
# we don't need as many filters
sed -i 's/filters=255/filters=21/g' cfg/yolov4-tiny-bus-trucks.cfg
sed -i 's/filters=57/filters=33/g' cfg/yolov4-tiny-bus-trucks.cfg
训练模型
现在我们准备好了，剩下的就是启动模型训练。
!./darknet/darknet detector train darknet/data/obj.data ./darknet/cfg/yolov4-tiny-bus-trucks.cfg yolov4-tiny.conv.29 -dont_show -mapLastAt
在这个情况下，这个训练大约花了一个小时。
现在，你可以在你的图像上运行预测，以获取类别和边界框。
from torch_snippets import Glob, stem, show, read
# upload your own images to a folder
image_paths = Glob('images-of-trucks-and-busses')
for f in image_paths:
!./darknet detector test \
data/obj.data cfg/yolov4-tiny-bus-trucks.cfg\
backup/yolov4-tiny-bus-trucks_4000.weights {f}
!mv predictions.jpg {stem(f)}_pred.jpg
for i in Glob('*_pred.jpg'):
show(read(i, 1), sz=20)
最后
使用YOLO并不复杂，正如我们所看到的。我们可以克隆并用于我们的用例的高效实现。
我们没有深入探讨这个算法的工作原理，因为希望在未来的文章中采用自上而下的方法更详细地介绍。