基于自定義數據集的YOLOv8模型實戰

作者：朱先忠 2024-10-25 08:30:57

本文將通過一個完整的實戰案例來展示使用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模型的輸出做點什么了。其實，你可以通過上述列舉的很多種方式將模型應用于自己的項目中。