用自己的数据集训练YOLOv8实例分割模型

YOLOv8 于 2023 年 1 月 10 日推出。截至目前，这是计算机视觉领域分类、检测和分割任务的最先进模型。该模型在准确性和执行时间方面都优于所有已知模型。

与之前的所有 YOLO 模型相比，ultralytics 团队在使该模型更易于使用方面做得非常好 — 你甚至不必再克隆 git 存储库！

01创建图像数据集

在这篇文章中，我创建了一个非常简单的示例，展示了在数据上训练 YOLOv8 所需执行的所有操作，特别是针对分割任务。数据集很小，并且模型“易于学习”，这样我们就可以在简单的 CPU 上训练几秒钟后得到令人满意的结果。

我们将创建一个黑色背景白色圆圈的数据集。圆圈的大小各不相同。我们将训练一个模型来分割图像内的圆圈。

数据集如下所示：

数据集是使用以下代码生成的：

import numpy as np
from PIL import Image
from skimage import draw
import random
from pathlib import Path

def create_image(path, img_size, min_radius):
    path.parent.mkdir( parents=True, exist_ok=True )
    
    arr = np.zeros((img_size, img_size)).astype(np.uint8)
    center_x = random.randint(min_radius, (img_size-min_radius))
    center_y = random.randint(min_radius, (img_size-min_radius))
    max_radius = min(center_x, center_y, img_size - center_x, img_size - center_y)
    radius = random.randint(min_radius, max_radius)

    row_indxs, column_idxs = draw.ellipse(center_x, center_y, radius, radius, shape=arr.shape)
    
    arr[row_indxs, column_idxs] = 255

    im = Image.fromarray(arr)
    im.save(path)

def create_images(data_root_path, train_num, val_num, test_num, img_size=640, min_radius=10):
    data_root_path = Path(data_root_path)
    
    for i in range(train_num):
        create_image(data_root_path / 'train' / 'images' / f'img_{i}.png', img_size, min_radius)
        
    for i in range(val_num):
        create_image(data_root_path / 'val' / 'images' / f'img_{i}.png', img_size, min_radius)
        
    for i in range(test_num):
        create_image(data_root_path / 'test' / 'images' / f'img_{i}.png', img_size, min_radius)

create_images('datasets', train_num=120, val_num=40, test_num=40, img_size=120, min_radius=10)

02创建标签

现在有了图像数据集，我们需要为图像创建标签。通常，需要为此做一些手动工作，但由于我们创建的数据集非常简单，因此创建生成标签的代码非常容易：

from rasterio import features

def create_label(image_path, label_path):
    arr = np.asarray(Image.open(image_path))

    # There may be a better way to do it, but this is what I have found so far
    cords = list(features.shapes(arr, mask=(arr >0)))[0][0]['coordinates'][0]
    label_line = '0 ' + ' '.join([f'{int(cord[0])/arr.shape[0]} {int(cord[1])/arr.shape[1]}' for cord in cords])

    label_path.parent.mkdir( parents=True, exist_ok=True )
    with label_path.open('w') as f:
        f.write(label_line)

for images_dir_path in [Path(f'datasets/{x}/images') for x in ['train', 'val', 'test']]:
    for img_path in images_dir_path.iterdir():
        label_path = img_path.parent.parent / 'labels' / f'{img_path.stem}.txt'
        label_line = create_label(img_path, label_path)

以下是标签文件内容的示例：

0 0.0767 0.08433 0.1417 0.08433 0.1417 0.0917 0.15843 0.0917 0.15843 0.1 0.1766 0.1 0.1766 0.10844 0.175 0.10844 0.175 0.1177 0.18432 0.1177 0.18432 0.14333 0.1918 0.14333 0.1918 0.20844 0.18432 0.20844 0.18432 0.225 0.175 0.225 0.175 0.24334 0.1766 0.24334 0.1766 0.2417 0.15843 0.2417 0.15843 0.25 0.1417 0.25 0.1417 0.25846 0.0767 0.25846 0.0767 0.25 0.05 0.25 0.05 0.2417 0.04174 0.2417 0.04174 0.24334 0.04333 0.24334 0.04333 0.225 0.025 0.225 0.025 0.20844 0.01766 0.20844 0.01766 0.14333 0.025 0.14333 0.025 0.1177 0.04333 0.1177 0.04333 0.10844 0.04174 0.10844 0.04174 0.1 0.05 0.1 0.05 0.0917 0.0767 0.0917 0.0767 0.08433

标签对应于该图像：

标签内容只是一个文本行。每张图像中只有一个对象（圆圈），每个对象在文件中由一行表示。如果每张图像中有多个对象，则应为每个标记的对象创建一条线。

第一个 0 代表标签的类别类型。因为我们只有一种类类型（圆形），所以总是 0。如果数据中有多个类，则应该将每个类映射到一个数字（0、1、2…），并在标签文件中使用该数字。

所有其他数字表示标记对象的边界多边形的坐标。格式为 并且坐标相对于图像的大小 - 您应该将坐标标准化为 1x1 图像大小。例如，如果有一个点 (15, 75) 并且图像大小为 120x120，则归一化点为 (15/120, 75/120) = (0.125, 0.625)。

在处理图像库时获取坐标的正确方向性总是令人困惑。所以为了明确这一点，对于YOLO来说，X坐标是从左到右，Y坐标是从上到下。

03YAML 配置

我们有了图像和标签。现在需要使用数据集配置创建一个 YAML 文件：

yaml_content = f'''
train: train/images
val: val/images
test: test/images

names: ['circle']
    '''
    
with Path('data.yaml').open('w') as f:
    f.write(yaml_content)

请注意，如果你有更多对象类类型，则需要按照在标签文件中的顺序将它们添加到名称数组中。第一个是 0，第二个是 1，等等...

04数据集文件结构

让我们看看使用 Linux 树命令创建的文件结构：

tree .

data.yaml
datasets/
├── test
│   ├── images
│   │   ├── img_0.png
│   │   ├── img_1.png
│   │   ├── img_2.png
│   │   ├── ...
│   └── labels
│       ├── img_0.txt
│       ├── img_1.txt
│       ├── img_2.txt
│       ├── ...
├── train
│   ├── images
│   │   ├── img_0.png
│   │   ├── img_1.png
│   │   ├── img_2.png
│   │   ├── ...
│   └── labels
│       ├── img_0.txt
│       ├── img_1.txt
│       ├── img_2.txt
│       ├── ...
|── val
|   ├── images
│   │   ├── img_0.png
│   │   ├── img_1.png
│   │   ├── img_2.png
│   │   ├── ...
|   └── labels
│       ├── img_0.txt
│       ├── img_1.txt
│       ├── img_2.txt
│       ├── ...

05训练模型

现在我们有了图像和标签，可以开始训练模型了。首先让我们安装这个包：

pip install ultralytics==8.0.38

ultralytics 库变化非常快，有时会破坏 API，所以我更喜欢坚持使用一个版本。下面的代码取决于版本 8.0.38（我写这篇文章时的最新版本）。如果你升级到较新的版本，也许需要进行一些代码调整才能使其正常工作。

开始训练：

from ultralytics import YOLO

model = YOLO("yolov8n-seg.pt")

results = model.train(
        batch=8,
        device="cpu",
        data="data.yaml",
        epochs=7,
        imgsz=120,
    )

为了这篇文章的简单性，我使用 Nano 模型 (yolov8n-seg)，我仅在 CPU 上训练它，只有 7 个 epoch。在我的笔记本电脑上进行训练只花了几秒钟。

有关可用于训练模型的参数的更多信息，你可以点击这里查看。

06理解结果

训练完成后，你将在输出末尾看到与此类似的一行：

Results saved to runs/segment/train60

让我们看一下此处找到的一些结果：

验证标注：

from IPython.display import Image as show_image
show_image(filename="runs/segment/train60/val_batch0_labels.jpg")

在这里我们可以看到验证集部分的真实标注。这应该几乎完全对齐。如果你发现这些标注没有很好地覆盖物体，那么标注很可能是不正确的。

预测的验证标签：

show_image(filename="runs/segment/train60/val_batch0_pred.jpg")

在这里，我们可以看到训练模型对部分验证集（与我们上面看到的相同部分）所做的预测。这可以让你了解模型的表现如何。请注意，为了创建此图像，应选择置信度阈值，此处使用的阈值是 0.5，这并不总是最佳阈值（我们将在稍后讨论）。

精度曲线：

要理解这张图表和下一张图表，你需要熟悉精确度和召回率概念。这里很好地解释了它们的工作原理。

show_image(filename="runs/segment/train60/MaskP_curve.png")

模型检测到的每个对象都有一定的置信度，通常，如果在声明“这是一个圆”时尽可能确定对你来说很重要，你将仅使用高置信度值（高置信度阈值）。当然，这需要权衡——你可能会错过一些“圈子”。另一方面，如果你想“捕获”尽可能多的“圆圈”，但要权衡其中一些不是真正的“圆圈”，可以同时使用低置信度值和高置信度值（低置信度阈值）。

上图（以及下图）可帮助你决定使用哪个置信阈值。在我们的例子中，我们可以看到，对于高于 0.128 的阈值，我们获得 100% 的精度，这意味着所有对象都被正确预测。

请注意，因为我们实际上是在做分割任务，所以我们需要担心另一个重要的阈值——IoU（交并集），如果你不熟悉，你可以在这里阅读介绍文章。对于此图表，使用的 IoU 为 0.5。

召回曲线：

show_image(filename="runs/segment/train60/MaskR_curve.png")

   

在这里你可以看到召回率图表，随着置信度阈值的上升，召回率下降。这意味着你“抓住”的“圆圈”更少。

在这里你可以明白为什么在这种情况下使用 0.5 置信度阈值是一个坏主意。对于 0.5 的阈值，可以获得大约 90% 的召回率。然而，在精度曲线中，我们看到对于高于 0.128 的阈值，我们获得 100% 的精度，因此我们不需要达到 0.5，我们可以安全地使用 0.128 阈值并获得 100% 的精度和几乎 100% 记起 ：）

精确率-召回率曲线：

这是精确率-召回率曲线的一个很好的解释

show_image(filename="runs/segment/train60/MaskPR_curve.png")

 

这里我们可以清楚地看到之前得出的结论，对于这个模型，我们可以达到几乎 100% 的精确率和 100% 的召回率。

这个图表的缺点是我们看不到应该使用什么阈值，这就是为什么我们仍然需要上面的图表。

随时间的推移损失

show_image(filename="runs/segment/train60/results.png")

 

在这里，你可以看到不同损失在训练过程中如何变化，以及它们在每个时期后在验证集上的表现如何。

关于损失，以及可以从这些图表中得出的结论，有很多话要说，但是，这超出了本文的范围。我只是想说你可以在这里找到它:)

07使用训练好的模型

可以在结果目录中找到的另一件事是模型本身。以下是如何在新图像上使用该模型：

my_model = YOLO('runs/segment/train60/weights/best.pt')
results = list(my_model('datasets/test/images/img_5.png', conf=0.128))
result = results[0]

结果列表可能有多个值，每个值对应一个检测到的对象。因为在此示例中，每个图像中只有一个对象，所以我们采用第一个列表项。

你可以看到我在这里传递了我们之前找到的最佳置信度阈值 (0.128)。

有两种方法可以获取检测到的对象在图像中的实际位置。选择正确的方法取决于你打算如何处理结果。我将展示这两种方式。

result.masks.segments

[array([[    0.10156,     0.34375],
        [    0.09375,     0.35156],
        [    0.09375,     0.35937],
        [   0.078125,       0.375],
        [   0.070312,       0.375],
        [     0.0625,     0.38281],
        [    0.38281,     0.71094],
        [    0.39062,     0.71094],
        [    0.39844,     0.70312],
        [    0.39844,     0.69531],
        [    0.41406,     0.67969],
        [    0.42187,     0.67969],
        [    0.44531,     0.46875],
        [    0.42969,     0.45312],
        [    0.42969,     0.41406],
        [    0.42187,     0.40625],
        [    0.41406,     0.40625],
        [    0.39844,     0.39062],
        [    0.39844,     0.38281],
        [    0.39062,       0.375],
        [    0.38281,       0.375],
        [    0.35156,     0.34375]], dtype=float32)]

这将返回对象的边界多边形，类似于我们传递标记数据的格式。

第二种方式：

result.masks.masks

tensor([[[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]]])

这将返回一个形状为 (1, 128, 128) 的张量，表示图像中的所有像素。作为对象一部分的像素接收 1，背景像素接收 0。

我们来看看面具是什么样子的：

import torchvision.transforms as T
T.ToPILImage()(result.masks.masks).show()

图像的预测分割

这是原始图像：

虽然并不完美，但对于很多应用来说已经足够好了，而且 IoU 肯定高于 0.5。

08结束语

总之，与之前的 Yolo 版本相比，新的 ultralytics 库更容易使用，尤其是对于分割任务，它现在是一等公民。你可以发现 Yolov5 也是 ultralytics 新软件包的一部分，所以如果你不想使用新的 Yolo 版本（它仍然是新的和实验性的），你可以使用众所周知的 yolov5：

转自新缸中之脑

审核编辑：汤梓红