Skip to main content

收集和标记图像以在 PyTorch 中训练 YOLOv5 目标检测模型

本教程将说明如何准备数据集以对自定义 YOLOv5 模型进行逐步训练。
Created on April 15|Last edited on April 19
/

欢迎来到我们 YOLOv5 教程系列的第二部分内容!如果你尚未了解本系列的第一部分内容,我建议你先读一下那部分内容,其中介绍了如何安装 YOLOv5 以在 Windows 和 Google Colab 中进行真实目标检测。在本报告中,我们将假定你已完成了这一工作。
接下来呢, 在完成了环境的安装工作后,你就可以开始着手准备数据集了。

章节



上手

我们每天都能看到计算机和机器人在自行学习,并且在不需要人为干预或明确编程的情况下进行持续改进。
不知你是否考虑过以下问题:
✓ 自动驾驶汽车是如何看见东西和安全驾驶的?
✓ Facebook 如何在图像中识别你并进行自动标记?
✓ 算法可以如何通过医学图像进行疾病诊断工作?
✓ 如何根据成熟度来对农作物进行分析和分类?
这些绝妙技术的核心在于计算机视觉。计算机视觉领域正在迅速发展。在某些情况下,它在解决很多视觉任务方面甚至已超过人类。
虽然计算机视觉领域涵盖了多种任务,但在本文中我们将着重介绍其中最普遍的的三种。
  1. 图像分类
  2. 图像分割
分类、本地化、检测以及分割之间的区别。
和上一个教程一样,今天我们会将注意力主要放在目标检测上。

什么是目标检测?

目标检测结合了分类和本地化的特点以确定某个对象是否存在于图像或视频中。
这是一种监督下机器学习模型,意味着我们需要为算法提供训练数据集,其中包含图像及其各自的标签。这听起来或许比较简单,但深度学习模型通常需要大量的数据才能完成训练。换言之,一个模型需要非常多的实例才能从未标记的图像中分辨出对象。
这一过程和人类学习的过程类似,重要的是提供既有质量又有数量的信息。数据的质量越高,训练的结果就越好。
Graph indicates how increasing amount of data increase the performance.

收集训练图像

在本教程中,我们的主要重点将是了解不同的数据收集方式以训练我们的自定义模型。
下面是一些可能的选项:

公开数据集


你可在这些地方轻松访问免费数据集。
  1. Google Dataset Search: 谷歌推出的这一搜索引擎可帮助研究人员轻松快速地访问数据集。其中包含超过 2500 万个数据集。
  2. Kaggle Datasets: Kaggle 可助力数据科学领域人士访问各类机器学习数据集。称得上是最适合本任务的资源。
  3. UCI Machine Learning Repository: UCI 页面建立于 1987 年,是世界上历史最长的免费数据集库。
  4. Visual Data: 顾名思义,这一搜索引擎中涵盖了专门用于计算机视觉任务的数据集。如果你正在寻找与分类、图像分割以及图像处理有关的数据集,那么这将是非常棒的资源。
  5. Papers With Code: 一个免费开源研究项目社区,其中包含代码与数据集。

数据增强

如果我们的图像受到局限,则可考虑使用数据增强技术。
此过程涉及通过新获取的外部数据来增强现有数据集。比如:裁剪、翻转、旋转、改变对比度和亮度等。
通过这一技术我们可以扩大现有小型训练数据的规模,因为从大多数目标检测模型的视角看来,其工作方式如同新图像一样。

免费在线图片

你可能听说过网上有很多图像。善于寻找资源的朋友们也可编写简单的爬虫程序来收集自己想要的数据集图像。

拍照

这是另一种可选项吗?拿起相机拍照吧。但是这样的做法很难一下子获得大量图像,不过你将能为特定任务收集到最合适的图像。

使用标记工具对数据进行标记

目前可用于对象检测的已标注数据量有限。这可能会是一个障碍。
如果你无法从上述提及的资源中找到已标注数据集,那么就需要自己动手去创建一个了。
GitHub 上提供了诸多免费工具,可帮助你对图像进行标注。
比较热门的几个是:
还有一些付费工具:
我将使用由OpenCV提供技术支持的开源库 OpenLabeling 但在开始标记数据之前,我们需要首先了解一下不同的边界框格式。
其中有:
每种格式都会使用其特定的边界框坐标表示形式。YOLOv5 和其他 YOLO 网络均会使用文件名称相同但扩展名会不同的两个文件。
一个文件是jpeg图像文件,另一个则是.txt文本文件(其中存储有关图像内标签的信息)。
行数表示图像中所存在的对象数量。每行有五个参数:
  • 对象类索引
  • 表示边界框中心位置的 (x,y) 坐标(x和y各为一个参数)
  • 边界框宽度
  • 边界框高度
坐标和边界框尺寸均按照图像尺寸的百分比被标准化为0至1之间的值
我将使用OpenLabeling工具来对图像进行标记。
我们现在就开始。

从GitHub获取标记库

首先,你需要从GitHub中下载并解压缩标记工具。

下载完成后提取文件夹。

需要进行标记的图像应放置在“images”文件夹中。

如前所述,YOLO格式中各图���的标签将被写入 bbox_txt 文件夹中的同名 .txt 文件。

要确保你已经完成运行修改后Open Labeling工具的requirements.txt文件中的必要库的安装。
只需打开文件夹位置即可,在地址栏中输入 cmd<enter>并输入:
pip install -r requirements.txt

为了启动工具并执行run.py文件,输入:
python run.py

其结果应该为:

顶部的滑动条可用于切换图像。

也可使用下方列出的快捷键进行操作

类别的数量在class.txt文件中进行了指定。在本文的例子中有4个类别:关上的门、打开的门、公共汽车和数字。

你在图像上绘制的所有边界框都会被自动添加至.txt文件中


将数据拆分为训练数据和验证数据:

标记完成后,我们就需要将数据拆分为训练集和验证集。
拆分的比例在train_test_split.py文件中进行了备注,你可以根据具体需求进行更改。
在命令提示符窗口中输入以下内容执行train_test_split.py:
python train_test_split.py
这将创建一个custom_dataset目录,将数据拆分到train和val 文件夹。

显示如下:

将自定义数据集文件夹复制进YOLOv5文件夹(在教程一中创建)。


为数据创建YAML文件:

创建YAML文件前,我们先来解答一下你可能会问的一个问题:

究竟什么是YAML文件呢?

YAML文件是指以YAML格式编写的文件(YAML 不是一种营销标记语言)。其通常会在配置文件中使用。

创建YAML文件

前往yolov5-master中的数据文件夹,新建一个名为custom_data.yaml的文件

最后对custom_data.yaml进行编辑。
更新train和val文件夹的路径:
train: custom_dataset/train/
val: custom_dataset/val/

# number of classes
nc: 4

# class names
names: ['closed_door', 'opened_door', 'bus', 'number']
形式如下:

记得保存!

使用W&B工件上传数据集

如果你还没有W&B账户,请花几分钟的时间注册一个免费的W&B账户
完成注册后即可快速创建一个W&B项目:

然后使用以下代码将W&B安装至YOLO中:
pip install wandb

接下来使用以下命令登录W&B:
wandb login
注意:如果出现错误:'wandb' 不被认为内部或外部命令,也不是可运行的程序或批处理文件。 这表明你可能没有将Scripts目录添加到Python的安装路径中。我们已在第一部分的这里对相关的步骤进行过讨论了。变通的方法是输入:python -m wandb login
现在是记录数据集的时候了。输入以下代码即可:
python utils/loggers/wandb/log_dataset.py --project YOLO --data data/custom_dataset.yaml
记得要把项目名称改成自己的项目名。
大功告成……

现在你就可以在W&B工件中找到这些数据集了。
在你的val.table中可轻松找到所有这些内容。

对于那些想与同事进行分享的人而言,十分便利的一点是,此表格和诸多其他功能可以交互式地嵌入到报告中。

id
train_image
Classes
name
1
2
3
一些人更愿意通过推文学习,还有些人则喜欢通过视频学习。如果你在此过程中遇到了任何麻烦,或是有任何问题,都可随时在下方评论区中留言。与此同时还可以看看这个由W&B团队成员 Ivan Goncharov 带来的精彩视频。s

视频讲解




推荐阅读


Iterate on AI agents and models faster. Try Weights & Biases today.
File<(table)>