基于自定义数据集的YOLOv8模型实战

译文 精选
人工智能
本文将通过一个完整的实战案例来展示使用Python、命令行或Google Colab等方式在自定义数据集上训练自己的计算机视觉模型。

译者 | 朱先忠

审校 | 重楼

简介

当前,Ultralytics公司研制的高级YOLOv8模型成为解决计算机视觉问题的最佳方法之一,同时该模型也最大限度地减少了有关开发过程遭遇的麻烦。YOLOv8是Ultralytics YOLO(You Only Look Once)系列模型的第8次也是最新一次迭代,与其他迭代一样,它使用卷积神经网络(CNN)来预测对象类别及其边界框。YOLO系列对象探测器以准确和快速而闻名,并提供了一个基于PyTorch构建的平台,从而简化了人们从头开始创建模型的大部分过程。

尤其重要的是,YOLOv8也是一个非常灵活的模型。它可以在各种平台上使用你选择的任何数据集进行训练,预测模型可以基于许多数据源运行。本文将作为一个全面的教程,涵盖训练和运行YOLOv8模型的许多不同方法,并展示每种方法的优缺点,这些方法将根据你的硬件和数据集帮你选择最合适的开发方案。

【注意】本文中创建上述示例数据集时使用的所有图像均由作者本人拍摄。

开发环境搭建

要开始训练我们的YOLOv8模型,第一步是决定我们想在什么样的环境中训练我们的模型(请记住,训练和运行模型是单独的任务)。

归纳来看,可供我们选择的环境在很大程度上可以分为两类:基于本地运行和基于云端方案运行。

通过基于本地的训练,我们基本上是使用设备的物理硬件直接在我们的系统上运行训练过程。在基于本地的训练中,YOLOv8为我们提供了两个选项:Python API和CLI。这两种选择的结果或速度没有真正的区别,因为同样的过程是在幕后进行的;唯一的区别在于训练的设置和运行方式。

另一方面,基于云端的训练允许你利用云服务器的硬件。通过使用互联网,你可以连接到云运行时并执行代码,就像在本地计算机上一样,只是现在它在云硬件上运行。

到目前为止,最受欢迎的机器学习云平台是谷歌Colab。它使用Jupyter笔记本格式,允许用户创建可以编写和运行代码片段的“单元格”,并提供与Google Drive和Github的强大集成。

你决定使用哪种环境在很大程度上取决于你可用的硬件。如果你有一个配备高端NVIDIA GPU的强大系统,基于本地的训练可能会对你很有效。如果你的本地机器的硬件不符合机器学习的配置建议,或者你只是想要比本地更多的计算能力,谷歌Colab可能是你更恰当的选择。

谷歌Colab最大的好处之一是它免费提供了一些计算资源,但也有一个简单的升级路径,可以让你利用更快的计算硬件。即使你已经有了一个强大的系统,如果谷歌Colab的高层计划中提供的更快的GPU比你现有的硬件有显著的性能提升,你也可以考虑使用谷歌Colab。使用免费计划,你只能使用NVIDIA T4,其性能大致相当于RTX 2070。对于更高级别配置的计划,L4(性能约为4090)和A100(性能约2 4090)是可用的。在比较GPU时,请记住VRAM的数量是机器学习性能的主要决定因素。

数据集

为了开始训练模型,你需要大量数据来训练它。对象检测数据集通常由各种对象的图像集合组成,此外还有一个围绕对象的“边界框”,指示其在图像中的位置。

检测到的对象周围的边界框示例

YOLOv8兼容的数据集具有特定的结构。它们主要分为有效文件夹、训练文件夹和测试文件夹,分别用于模型的验证、训练和测试(验证和测试之间的区别在于,在验证过程中,结果用于调整模型以提高其准确性;而在测试过程中,结果仅用于提供模型真实世界准确性的衡量标准)。

在每个文件夹中,数据集进一步分为两个文件夹:图像(images)和标签(labels)文件夹。这两个文件夹的内容紧密相连。

顾名思义,images文件夹包含数据集的所有对象图像。这些图像通常具有方形纵横比、低分辨率和小的文件尺寸。

labels文件夹包含边界框在每个图像中的位置和大小的数据,以及每个图像表示的对象的类型(或类别)。例如:

这里的第一行数据代表图像中存在的单个对象。在每一行中,第一个数字表示对象的类别,第二个和第三个数字表示边界框中心的x和y坐标,第四个和第五个数字表示边缘框的宽度和高度。

images和labels文件夹中的数据通过文件名链接在一起。images文件夹中的每个图像在labels文件夹中都有一个具有相同文件名的相应文件;反之亦然。在数据集中,images和labels文件夹中总是有匹配的文件对,它们具有相同的文件名,但文件扩展名不同;.jpg用于图像,.txt用于标签。.jpg图片中每个对象的边界框数据包含在相应的.txt文件中。

YOLOv8兼容数据集的典型文件结构。来源:Ultralytics公司YOLO文档

有几种方法可以获得与YOLOv8兼容的数据集来开始训练模型。你可以创建自己的数据集,也可以使用互联网上预先配置的数据集。在本教程中,我们将使用CVAT创建自己的数据集,并使用Kaggle查找预配置的数据集。

CVAT

CVAT(CVAT.ai)是一个注释工具,它允许你通过手动方式为图像和视频添加标签来创建自己的数据集。

创建帐户并登录后,开始注释的过程很简单。只需创建一个项目,给它一个合适的名称,并根据需要为尽可能多的对象类型/类别添加标签。

作者在cvat.ai.Video上创建新项目和标签

创建一个新任务,并上传你希望成为数据集一部分的所有图像。点击“Submit & Open”按钮,项目下应创建一个新任务,其中包含一个作业。

在cvat.ai上创建新任务和作业

打开此作业将允许你启动注释过程。你可以使用矩形工具为数据集中的每个图像创建边界框和标签。

使用cvat.ai上的矩形工具创建边界框

在为所有图像添加注释后,返回任务并选择“Actions(动作)”→“Export task dataset(导出任务数据集)”,然后选择YOLOv8 Detection 1.0作为导出格式。下载任务数据集后,你会发现它只包含labels文件夹,而不包含images文件夹(除非你在导出时选择了“Save images(保存图像)”选项)。你必须手动创建images文件夹并将图像移动到那里(你可能希望首先将图像压缩到较低的分辨率,例如640x640)。记住不要更改文件名,因为它们必须与labels文件夹中.txt文件的文件名匹配。你还需要决定如何在验证、训练和测试之间分配图像(训练是其中最重要的)。

作者从cvat.ai.Image导出的示例数据集

至此,你的数据集已完成并准备好使用了!

Kaggle

Kaggle(Kaggle.com)是最大的在线数据科学社区之一,也是探索数据集的最佳网站之一。你可以通过简单地搜索他们的网站来尝试找到你需要的数据集,除非你正在寻找非常具体的东西;否则,你很可能会找到意向的数据集。然而,Kaggle上的许多数据集不是YOLOv8兼容的格式和/或与计算机视觉无关;因此,你可能想在查询中包含“YOLOv9”来优化你的搜索。

你可以通过数据集的Data Explorer(页面右侧)中的文件结构来判断数据集是否与YOLOv8兼容。

Kaggle上兼容YOLOv8的数据集示例

如果数据集相对较小(几MB)和/或你想在本地训练,那么可以直接从Kaggle下载数据集。但是,如果你计划在Google Colab上使用大型数据集进行训练,最好从笔记本文件本身进行数据集检索(更多信息见下文)。

训练模型

训练过程将根据你是在本地还是在云端进行训练而有所不同。

本地

为所有训练文件创建一个项目文件夹。在本教程中,我们将称之为yolov8-project。然后,将数据集移动/复制到此文件夹下。

接下来,使用所需的YOLOv8依赖项设置Python虚拟环境:

python3 -m venv venv
source venv/bin/activate
pip3 install ultralytics

然后,创建一个名为config.yaml的配置文件。这个文件将会指定用于训练的重要数据集信息:

path: /Users/oliverma/yolov8-project/dataset/ # absolute path to dataset
test: test/images # relative path to test images
train: train/images # relative path to training images
val: val/images # relative path to validation images

# classes
names:
0: bottle

在上面配置信息的path部分,提供的是指向数据集根目录的绝对文件路径。你也可以使用相对文件路径,但这取决于config.yaml的相对位置。

然后,在test、train和val部分,提供用于测试、训练和验证的图像的位置(如果你只有训练图像,则只需对所有这三种操作均使用train/images)。

在names部分,指定每个类别的名称。这些信息通常可以在任何YOLOv8数据集的data.yaml文件中找到。

如前所述,Python API或CLI(命令行方式)都可以用来进行本地训练。

Python API方式

创建另一个名为main.py的文件。这是实际训练开始的地方:

from ultralytics import YOLO
model = YOLO("yolov8n.yaml")
model.train(data="config.yaml", epochs=100)

通过将我们的模型初始化为YOLO("yolov8n.yaml"),我们基本上是从头开始创建一个新模型。我们使用yolov8n是因为它是最快的模型,但根据你自己的使用情况,你也可以选择使用其他模型。

性能指标

YOLOv8变体的性能指标。来源:Ultralytics YOLO文档

最后,我们开始训练模型,并传递配置文件和迭代次数,或训练轮数。一个比较好的训练指标是使用300个训练轮数,但你可能想根据数据集的大小和硬件的速度来调整这个数字。

你可能还希望包括一些更有用的设置:

  • imgsz:将所有图像调整到指定的大小。例如,imgsz=640会将所有图像的大小调整为640x640。如果你创建了自己的数据集并且没有调整图像大小,这将非常有用。
  • device:指定要在哪个设备上训练。默认情况下,YOLOv8会尝试在GPU上训练,并使用CPU训练作为后备,但如果你在M系列Mac上训练,你必须使用device="mps"以便使用苹果电脑上的Metal Performance Shaders(mps)后端进行GPU加速训练。
  • 有关所有训练参数的更多信息,请访问https://docs.ultralytics.com/modes/train/#train-settings。

你的项目目录现在应该看起来类似于下面的样子:

项目目录的示例文件结构

现在,我们终于准备好开始训练我们的模型了。为此,只需要在项目目录中打开一个终端并运行:

python3 main.py

随着训练的进行,终端将显示每个训练世代的训练进度信息。

终端中显示的每个训练世代的训练进度

训练结果将保存在路径runs/detect/train(或者train2,train3,等)。注意,这里包括了权重数据(文件扩展名为.pt),这对以后运行模型很重要;还有文件results.png,它显示了许多包含相关训练统计数据的图表。

results.png文件中的示例图

CLI方式

在项目目录中打开一个新终端并运行以下命令:

yolo detect train data=config.yaml model=yolov8n.yaml epochs=100

此命令可以使用上面为Python API列出的相同参数进行修改。例如:

yolo detect train data=config.yaml model=yolov8n.yaml epochs=300 imgsz=640 device=mps

训练将开始,进度将显示在终端上。其余的训练过程与Python CLI相同。

谷歌Colab方式

导航到https://colab.research.google.com/,并为训练创建一个新的笔记本文件。

在训练之前,请确保通过选择右上角的“Change runtime type(更改运行时类型)”连接到GPU运行时。CPU运行时的训练将非常缓慢。

将笔记本电脑运行时从CPU更改为T4 GPU

在开始使用Google Colab进行任何训练之前,我们首先需要将数据集导入笔记本文件。直观地说,最简单的方法是将数据集上传到谷歌云端硬盘,并从那里导入到我们的笔记本文件中。然而,上传任何大于几MB的数据集都需要非常长的时间。解决方法是将数据集上传到远程文件托管服务(如Amazon S3甚至是Kaggle),并将数据集直接从那里拉入我们的Colab笔记本文件。

从Kaggle导入

以下是如何将Kaggle数据集直接导入Colab笔记本文件的说明:

在Kaggle账户设置中,向下滚动到API并选择“Create New Token”命令。这将下载一个名为kaggle.json的文件。

在笔记本单元格中运行以下命令:

!pip install kaggle
from google.colab import files
files.upload()

上传刚刚下载的kaggle.json文件,然后运行以下命令:

!mkdir ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json
!kaggle datasets download -d [DATASET] # replace [DATASET] with the desired dataset ref

数据集将作为zip存档下载。只需要使用unzip命令提取有关内容:

!unzip dataset.zip -d dataset

开始训练

在笔记本文件的文件资源管理器中创建一个新的config.yaml文件,并按照前面的描述进行配置。Colab笔记本文件中的默认工作目录是/content/,因此数据集的绝对路径将是/content/[dataset folder]。例如:

path: /content/dataset/ # absolute path to dataset
test: test/images # relative path to test images
train: train/images # relative path to training images
val: val/images # relative path to validation images

# classes
names:
0: bottle

确保检查一下数据集的文件结构,以确保config.yaml中指定的路径准确。有时数据集会被放置在多个级别的文件夹中。

然后,将以下内容作为单元格运行:

!pip install ultralytics

import os
from ultralytics import YOLOmodel = YOLO("yolov8n.yaml")
results = model.train(data="config.yaml", epochs=100)

前面提到的用于修改本地训练设置的参数也适用于这里。

与本地训练类似,结果、权重和图表将保存在runs/detect/train中。

在本地进行预测

无论你是在本地还是在云端进行训练,预测都必须在本地运行。

在模型完成训练后,runs/detect/train/weights中会有两个权重,分别命名为best.pt和last.pt,分别是最佳训练轮次和最新训练轮次的权重。在本教程中,我们将使用best.pt运行模型。

如果你在本地进行训练,请将best.pt移动到方便的位置(例如我们的项目文件夹yolov8-project)以运行预测。如果你在云端训练,请将best.pt下载到你的本地设备上。在Google Colab上,右键单击笔记本资源管理器中的文件,然后选择“Download”(下载)命令。

在谷歌Colab上下载权重

与本地训练类似,预测可以通过Python API或CLI运行。

使用Python API预测

在与best.pt相同的位置,创建一个名为predict.py的新文件:

from ultralytics import YOLO
model = YOLO("best.pt")
results = model(source=0, show=True, conf=0.25, save=True)

与训练类似,存在许多有用的参数可以修改预测设置:

  • source:控制预测的输入源。source=0将网络摄像头设置为输入源。更多信息请见下文。
  • show:如果为True,则在屏幕上显示预测、边界框和置信度。
  • conf:要考虑预测的最小置信度阈值。
  • save:如果为True,则将预测结果保存到run/retect/repredict(或predict2、predict3等)。
  • device:如前所述,在M系列Mac上使用device=“mps”。

有关预测参数的完整列表,请访问:

https://docs.ultralytics.com/modes/predict/#inference-arguments

使用CLI预测

运行以下CLI命令即可启动模型:

python3 predict.py

通过实时网络摄像头馈送运行YOLOv8模型预测

CLI命令

使用CLI方式进行预测应用的命令是:

yolo detect predict model=best.pt source=0 show=True cnotallow=0.25 save=True

注意到,这里所使用的参数与Python API中的参数相同。

将YOLOv8模型集成到实际场景

现在,我们已经能够在实时网络摄像头上成功地运行起我们的模型,但这又会怎样呢?我们如何实际使用这个模型并将其集成到实战项目中?

让我们从输入和输出的角度来考虑这个问题。为了使这个模型成功地应用于我们构建的外部应用程序中,此模型必须能够接受有用的输入并产生有用的输出。值得庆幸的是,YOLOv8模型的灵活性使得将模型集成到各种使用场景中成为可能。

例如,我们可以使用source=0将网络摄像头设置为预测的输入源。然而,YOLOv8模型可以利用比这更多的输入源。以下是几个例子:

results = model(source="path/to/image.jpg", show=True, conf=0.25, save=True) # static image
results = model(source="screen", show=True, conf=0.25, save=True) # screenshot of current screen
results = model(source="https://ultralytics.com/images/bus.jpg", show=True, conf=0.25, save=True) # image or video URL
results = model(source="path/to/file.csv", show=True, conf=0.25, save=True) # CSV file
results = model(source="path/to/video.mp4", show=True, conf=0.25, save=True) # video file
results = model(source="path/to/dir", show=True, conf=0.25, save=True) # all images and videos within directory
results = model(source="path/to/dir/**/*.jpg", show=True, conf=0.25, save=True) # glob expression
results = model(source="https://www.youtube.com/watch?v=dQw4w9WgXcQ", show=True, conf=0.25, save=True) # YouTube video URL

有关预测源和输入选项的完整列表,请访问:

https://docs.ultralytics.com/modes/predict/#inference-sources。

每当我们运行预测时,YOLOv8都会以Results对象列表的形式返回大量有价值的数据,其中包括有关预测的边界框、分割掩码、关键点、类别概率和定向包围盒(OBB)的信息。

由于我们在代码中将预测结果分配给results变量,因此我们可以使用它来检索有关预测的信息:

from ultralytics import YOLO

model = YOLO("best.pt")

results = model(source="bottles.jpg", show=True, conf=0.25, save=True)

print("Bounding boxes of all detected objects in xyxy format:")
for r in results:
print(r.boxes.xyxy)

print("Confidence values of all detected objects:")
for r in results:
print(r.boxes.conf)

print("Class values of all detected objects:")
for r in results:
print(r.boxes.cls)

本教程中包含的输出结果类型太多,但你可以通过访问链接以了解更多信息:

https://docs.ultralytics.com/modes/predict/#working-results。

这里提供的只是一个非常基本的例子,说明你可以用YOLOv8模型的输出做点什么了。其实,你可以通过上述列举的很多种方式将模型应用于自己的项目中。

结论

在本文中,我们终于实现了从头开始制作我们自己的YOLOv8兼容数据集,从Kaggle导入数据集,使用包括Python API、CLI和Google Colab在内的多种环境来训练模型,然后在本地运行我们的模型,并找到许多输入/输出方法,使我们能够在自己的项目中利用训练出的YOLOv8模型。

请记住,本教程的目的是作为YOLOv8或计算机视觉的学习起点。我们当前几乎没有触及YOLOv8模型错综复杂的表面,随着你对YOLOv9和计算机视觉越来越有经验,深入了解这个模型绝对是明智的。

话虽如此,如果你遵循了本教程并坚持到最后,那仍然是一个很大的成功。我希望这篇文章能帮助你对机器学习、计算机视觉和YOLOv8模型有一个基本的了解。也许你甚至对这个主题产生了热情,并将在未来继续学习更高级的内容。

译者介绍

朱先忠,51CTO社区编辑,51CTO专家博客、讲师,潍坊一所高校计算机教师,自由编程界老兵一枚。

原文标题:The Comprehensive Guide to Training and Running YOLOv8 Models on Custom Datasets,作者:Oliver Ma

责任编辑:姜华 来源: 51CTO内容精选
相关推荐

2024-10-16 16:49:44

定向边界框目标检测YOLOv8

2024-07-01 12:55:50

2023-01-12 12:20:29

YOLOv8模型

2024-07-22 13:49:38

YOLOv8目标检测开发

2024-07-09 08:50:23

2020-10-05 21:57:17

GitHub 开源开发

2024-07-11 08:25:34

2024-05-15 09:16:05

2023-09-12 13:59:41

OpenAI数据集

2024-11-06 16:56:51

2023-02-02 09:00:00

检测模型YOLOv8存储

2024-01-29 09:29:02

计算机视觉模型

2016-08-18 13:56:33

AndroidExecutorsubmit

2023-12-11 10:18:38

YOLOv8检测器实战

2024-10-07 11:12:55

2022-06-24 07:08:24

OHOS自定义服务

2024-11-18 17:31:27

2024-09-09 16:35:10

YOLO模型

2015-02-12 15:33:43

微信SDK

2024-10-14 17:43:05

点赞
收藏

51CTO技术栈公众号