用自己的數(shù)據(jù)集訓(xùn)練YOLOv8實(shí)例分割模型

YOLOv8 于 2023 年 1 月 10 日推出。截至目前，這是計(jì)算機(jī)視覺(jué)領(lǐng)域分類(lèi)、檢測(cè)和分割任務(wù)的最先進(jìn)模型。該模型在準(zhǔn)確性和執(zhí)行時(shí)間方面都優(yōu)于所有已知模型。

與之前的所有 YOLO 模型相比，ultralytics 團(tuán)隊(duì)在使該模型更易于使用方面做得非常好 — 你甚至不必再克隆 git 存儲(chǔ)庫(kù)！

01創(chuàng)建圖像數(shù)據(jù)集

在這篇文章中，我創(chuàng)建了一個(gè)非常簡(jiǎn)單的示例，展示了在數(shù)據(jù)上訓(xùn)練 YOLOv8 所需執(zhí)行的所有操作，特別是針對(duì)分割任務(wù)。數(shù)據(jù)集很小，并且模型“易于學(xué)習(xí)”，這樣我們就可以在簡(jiǎn)單的 CPU 上訓(xùn)練幾秒鐘后得到令人滿(mǎn)意的結(jié)果。

我們將創(chuàng)建一個(gè)黑色背景白色圓圈的數(shù)據(jù)集。圓圈的大小各不相同。我們將訓(xùn)練一個(gè)模型來(lái)分割圖像內(nèi)的圓圈。

數(shù)據(jù)集如下所示：

數(shù)據(jù)集是使用以下代碼生成的：

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創(chuàng)建標(biāo)簽

現(xiàn)在有了圖像數(shù)據(jù)集，我們需要為圖像創(chuàng)建標(biāo)簽。通常，需要為此做一些手動(dòng)工作，但由于我們創(chuàng)建的數(shù)據(jù)集非常簡(jiǎn)單，因此創(chuàng)建生成標(biāo)簽的代碼非常容易：

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)

以下是標(biāo)簽文件內(nèi)容的示例：

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

標(biāo)簽對(duì)應(yīng)于該圖像：

標(biāo)簽內(nèi)容只是一個(gè)文本行。每張圖像中只有一個(gè)對(duì)象（圓圈），每個(gè)對(duì)象在文件中由一行表示。如果每張圖像中有多個(gè)對(duì)象，則應(yīng)為每個(gè)標(biāo)記的對(duì)象創(chuàng)建一條線(xiàn)。

第一個(gè) 0 代表標(biāo)簽的類(lèi)別類(lèi)型。因?yàn)槲覀冎挥幸环N類(lèi)類(lèi)型（圓形），所以總是 0。如果數(shù)據(jù)中有多個(gè)類(lèi)，則應(yīng)該將每個(gè)類(lèi)映射到一個(gè)數(shù)字（0、1、2…），并在標(biāo)簽文件中使用該數(shù)字。

所有其他數(shù)字表示標(biāo)記對(duì)象的邊界多邊形的坐標(biāo)。格式為 并且坐標(biāo)相對(duì)于圖像的大小 - 您應(yīng)該將坐標(biāo)標(biāo)準(zhǔn)化為 1x1 圖像大小。例如，如果有一個(gè)點(diǎn) (15, 75) 并且圖像大小為 120x120，則歸一化點(diǎn)為 (15/120, 75/120) = (0.125, 0.625)。

在處理圖像庫(kù)時(shí)獲取坐標(biāo)的正確方向性總是令人困惑。所以為了明確這一點(diǎn)，對(duì)于YOLO來(lái)說(shuō)，X坐標(biāo)是從左到右，Y坐標(biāo)是從上到下。

03YAML 配置

我們有了圖像和標(biāo)簽。現(xiàn)在需要使用數(shù)據(jù)集配置創(chuàng)建一個(gè) 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)

請(qǐng)注意，如果你有更多對(duì)象類(lèi)類(lèi)型，則需要按照在標(biāo)簽文件中的順序?qū)⑺鼈兲砑拥矫Q(chēng)數(shù)組中。第一個(gè)是 0，第二個(gè)是 1，等等...

04數(shù)據(jù)集文件結(jié)構(gòu)

讓我們看看使用 Linux 樹(shù)命令創(chuàng)建的文件結(jié)構(gòu)：

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訓(xùn)練模型

現(xiàn)在我們有了圖像和標(biāo)簽，可以開(kāi)始訓(xùn)練模型了。首先讓我們安裝這個(gè)包：

pip install ultralytics==8.0.38

ultralytics 庫(kù)變化非?？欤袝r(shí)會(huì)破壞 API，所以我更喜歡堅(jiān)持使用一個(gè)版本。下面的代碼取決于版本 8.0.38（我寫(xiě)這篇文章時(shí)的最新版本）。如果你升級(jí)到較新的版本，也許需要進(jìn)行一些代碼調(diào)整才能使其正常工作。

開(kāi)始訓(xùn)練：

from ultralytics import YOLO

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

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

為了這篇文章的簡(jiǎn)單性，我使用 Nano 模型 (yolov8n-seg)，我僅在 CPU 上訓(xùn)練它，只有 7 個(gè) epoch。在我的筆記本電腦上進(jìn)行訓(xùn)練只花了幾秒鐘。

有關(guān)可用于訓(xùn)練模型的參數(shù)的更多信息，你可以點(diǎn)擊這里查看。

06理解結(jié)果

訓(xùn)練完成后，你將在輸出末尾看到與此類(lèi)似的一行：

Results saved to runs/segment/train60

讓我們看一下此處找到的一些結(jié)果：

驗(yàn)證標(biāo)注：

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

在這里我們可以看到驗(yàn)證集部分的真實(shí)標(biāo)注。這應(yīng)該幾乎完全對(duì)齊。如果你發(fā)現(xiàn)這些標(biāo)注沒(méi)有很好地覆蓋物體，那么標(biāo)注很可能是不正確的。

預(yù)測(cè)的驗(yàn)證標(biāo)簽：

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

在這里，我們可以看到訓(xùn)練模型對(duì)部分驗(yàn)證集（與我們上面看到的相同部分）所做的預(yù)測(cè)。這可以讓你了解模型的表現(xiàn)如何。請(qǐng)注意，為了創(chuàng)建此圖像，應(yīng)選擇置信度閾值，此處使用的閾值是 0.5，這并不總是最佳閾值（我們將在稍后討論）。

精度曲線(xiàn)：

要理解這張圖表和下一張圖表，你需要熟悉精確度和召回率概念。這里很好地解釋了它們的工作原理。

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

模型檢測(cè)到的每個(gè)對(duì)象都有一定的置信度，通常，如果在聲明“這是一個(gè)圓”時(shí)盡可能確定對(duì)你來(lái)說(shuō)很重要，你將僅使用高置信度值（高置信度閾值）。當(dāng)然，這需要權(quán)衡——你可能會(huì)錯(cuò)過(guò)一些“圈子”。另一方面，如果你想“捕獲”盡可能多的“圓圈”，但要權(quán)衡其中一些不是真正的“圓圈”，可以同時(shí)使用低置信度值和高置信度值（低置信度閾值）。

上圖（以及下圖）可幫助你決定使用哪個(gè)置信閾值。在我們的例子中，我們可以看到，對(duì)于高于 0.128 的閾值，我們獲得 100% 的精度，這意味著所有對(duì)象都被正確預(yù)測(cè)。

請(qǐng)注意，因?yàn)槲覀儗?shí)際上是在做分割任務(wù)，所以我們需要擔(dān)心另一個(gè)重要的閾值——IoU（交并集），如果你不熟悉，你可以在這里閱讀介紹文章。對(duì)于此圖表，使用的 IoU 為 0.5。

召回曲線(xiàn)：

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

在這里你可以看到召回率圖表，隨著置信度閾值的上升，召回率下降。這意味著你“抓住”的“圓圈”更少。

在這里你可以明白為什么在這種情況下使用 0.5 置信度閾值是一個(gè)壞主意。對(duì)于 0.5 的閾值，可以獲得大約 90% 的召回率。然而，在精度曲線(xiàn)中，我們看到對(duì)于高于 0.128 的閾值，我們獲得 100% 的精度，因此我們不需要達(dá)到 0.5，我們可以安全地使用 0.128 閾值并獲得 100% 的精度和幾乎 100% 記起 ：）

精確率-召回率曲線(xiàn)：

這是精確率-召回率曲線(xiàn)的一個(gè)很好的解釋

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

這里我們可以清楚地看到之前得出的結(jié)論，對(duì)于這個(gè)模型，我們可以達(dá)到幾乎 100% 的精確率和 100% 的召回率。

這個(gè)圖表的缺點(diǎn)是我們看不到應(yīng)該使用什么閾值，這就是為什么我們?nèi)匀恍枰厦娴膱D表。

隨時(shí)間的推移損失

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

在這里，你可以看到不同損失在訓(xùn)練過(guò)程中如何變化，以及它們?cè)诿總€(gè)時(shí)期后在驗(yàn)證集上的表現(xiàn)如何。

關(guān)于損失，以及可以從這些圖表中得出的結(jié)論，有很多話(huà)要說(shuō)，但是，這超出了本文的范圍。我只是想說(shuō)你可以在這里找到它:)

07使用訓(xùn)練好的模型

可以在結(jié)果目錄中找到的另一件事是模型本身。以下是如何在新圖像上使用該模型：

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]

結(jié)果列表可能有多個(gè)值，每個(gè)值對(duì)應(yīng)一個(gè)檢測(cè)到的對(duì)象。因?yàn)樵诖耸纠?，每個(gè)圖像中只有一個(gè)對(duì)象，所以我們采用第一個(gè)列表項(xiàng)。

你可以看到我在這里傳遞了我們之前找到的最佳置信度閾值 (0.128)。

有兩種方法可以獲取檢測(cè)到的對(duì)象在圖像中的實(shí)際位置。選擇正確的方法取決于你打算如何處理結(jié)果。我將展示這兩種方式。

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)]

這將返回對(duì)象的邊界多邊形，類(lèi)似于我們傳遞標(biāo)記數(shù)據(jù)的格式。

第二種方式：

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.]]])

這將返回一個(gè)形狀為 (1, 128, 128) 的張量，表示圖像中的所有像素。作為對(duì)象一部分的像素接收 1，背景像素接收 0。

我們來(lái)看看面具是什么樣子的：

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

圖像的預(yù)測(cè)分割

這是原始圖像：

雖然并不完美，但對(duì)于很多應(yīng)用來(lái)說(shuō)已經(jīng)足夠好了，而且 IoU 肯定高于 0.5。

08結(jié)束語(yǔ)

總之，與之前的 Yolo 版本相比，新的 ultralytics 庫(kù)更容易使用，尤其是對(duì)于分割任務(wù)，它現(xiàn)在是一等公民。你可以發(fā)現(xiàn) Yolov5 也是 ultralytics 新軟件包的一部分，所以如果你不想使用新的 Yolo 版本（它仍然是新的和實(shí)驗(yàn)性的），你可以使用眾所周知的 yolov5：

轉(zhuǎn)自新缸中之腦

審核編輯：湯梓紅