主页
avatar

Kared

YOLO V8 实战：私有数据集训练全指南

一点字  0分钟 Coding

1、训练环境配置

为了运行 YOLOv8，首先需要配置适当的环境。这里建议利用 Anaconda 来创建一个虚拟环境，这样有助于隔离和管理项目依赖，避免潜在的版本冲突。

在虚拟环境中安装 PyTorch 框架，并安装 Ultralytics 库，YOLOv8 的核心算法和功能都包含在这个库中。为了方便国内用户快速安装，我们将使用清华大学提供的 Python 软件包索引镜像站。运行以下命令来安装 Ultralytics 库：

pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

2、准备训练数据

在使用 YOLOv8 进行目标检测训练的过程中，采用 YOLO 数据格式通常是常见的选择。为了确保您的自定义数据集能够无缝对接 YOLOv8 的训练流程，接下来，我们将详细介绍如何将您的数据集转换为兼容的 YOLO 格式，确保数据集的结构和注释信息符合 YOLOv8 对输入数据的规范要求。

数据集组织结构

我们数据集被组织在一个名为 mydata 的文件夹（名称自定义）中，其目录结构包含了利用标注的 XML 文件和对应的图片，以及为 YOLOv8 训练准备的 TXT 格式标签数据文件：

  • mydata:
    • data : 存储训练集、测试集、验证集的文件列表
    • images : 存储数据集中的图片，是目标检测模型训练的视觉输入部分
    • labels : 存储与图片对应的 YOLO 数据格式的 TXT 标签数据文件
    • xml : 存储与图片对应的 XML 标注文件（可选，用来生成 labels 数据）
    • mask : 存储与图片对应的 MASK 标注文件（可选，用来生成 labels 数据）

将标注数据转化为 YOLO 数据格式

YOLO 数据格式是 YOLO 目标检测算法的图像标注格式，其注释文件以 .txt 结尾，包含类别索引和边界框（bounding box）的中心坐标、宽度和高度，这些值都被归一化到 [0,1]。

其文件内容的每一行代表图像中的一个目标，用空格隔开上述五个值（类别索引、x_center、y_center、width、height），如下所示：

0 0.50625 0.37 0.245 0.54
1 0.71625 0.47 0.16 0.24

对于使用 VOC 数据格式的数据集，标注信息是通过 XML 文件提供的，每个 XML 文件对应一张图片，并且包含了图像中所有目标的具体位置和类别等信息。您可以参考以下代码将 VOC 格式的 XML 文件转换为 YOLO 数据格式的 TXT 标签文件：

import os
import xml.etree.ElementTree as ET

class_mapping = {
    'cat': 0,
    'dog': 1,
}

def convert_bbox(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = ((box[0] + box[1]) / 2.0) * dw
    y = ((box[2] + box[3]) / 2.0) * dw
    w = (box[1] - box[0]) * dh
    h = (box[3] - box[2]) * dh
    return (x, y, w, h)

def voc2yolo(xml_label_path, yolo_txt_path):
    os.makedirs(yolo_txt_path, exist_ok=True)
  
    for filename in os.listdir(xml_label_path):
        if filename.endswith('.xml'):
            filepath = os.path.join(xml_label_path, filename)
            txt_name = os.path.splitext(filename)[0] + '.txt'
            txt_path = os.path.join(yolo_txt_path, txt_name)
  
            tree = ET.parse(filepath)
            root = tree.getroot()
            size = root.find('size')
            width = int(size.find('width').text)
            height = int(size.find('height').text)
  
            with open(txt_path, 'w') as f:
                for obj in root.iter('object'):
                    class_name = obj.find('name').text
                    if class_name in class_mapping:
                        class_id = class_mapping[class_name]

                        xmlbox = obj.find('bndbox')
                        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
                             float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
                        bbox = convert_bbox((width, height), b)
  
                        print(f"Writing bounding box to {txt_path}")
                        f.write(str(class_id) + " " + " ".join([str(a) for a in bbox]) + '\n')

src_directory = "/home/detect/datasets/mydata/"
images_path = os.path.join(src_directory, 'images')
labels_path = os.path.join(src_directory, 'labels')
xml_path = os.path.join(src_directory, 'xml')
voc2yolo(xml_path, labels_path)

对于图像分割任务的数据集，标注信息将以 png 文件的形式提供，每个 png 文件对应一张图片，包含了图像上的分割掩码（mask）信息，例如目标区域、类别等信息。您可以参考以下代码将 mask 转换为 YOLO 数据格式的 TXT 标签文件：

import os
import cv2

def mask2yolo(mask_label_path, yolo_txt_path):
    os.makedirs(yolo_txt_path, exist_ok=True)

    for filename in os.listdir(mask_label_path):
        if filename.endswith('.png'):
            filepath = os.path.join(mask_label_path, filename)
            txt_name = os.path.splitext(filename)[0] + '.txt'
            txt_path = os.path.join(yolo_txt_path, txt_name)
  
            image = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
            _, thresh = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY)
            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  
            with open(txt_path, 'w') as f:
                for contour in contours:
                    x, y, width, height = cv2.boundingRect(contour)
  
                    x_center = (x + width / 2) / image.shape[1]
                    y_center = (y + height / 2) / image.shape[0]
                    width_norm = width / image.shape[1]
                    height_norm = height / image.shape[0]
  
                    print(f"Writing bounding box to {txt_path}")
                    f.write(f"0 {x_center:.6f} {y_center:.6f} {width_norm:.6f} {height_norm:.6f}\n")

src_directory = "/home/detect/datasets/mydata/"
images_path = os.path.join(src_directory, 'images')
labels_path = os.path.join(src_directory, 'labels')
mask_path = os.path.join(src_directory, 'mask')
mask2yolo(mask_path, labels_path)

创建训练集和验证集

在准备用于YOLO模型训练的数据集时，我们通常需要分别创建包含训练集 (train) 和验证集 (val) 图像路径的文本文件。下面的代码段是自动执行此任务的脚本，它将查找指定图像目录下所有图像路径，然后将这些图像按照比例随机分配到训练集和验证集中，并生成训练集和验证集的划分文件。

import os
import random

train_ratio = 0.9
image_dir = '/home/detect/datasets/mydata/images'
train_txt_path = '/home/detect/datasets/mydata/data/train.txt'
val_txt_path = '/home/detect/datasets/mydata/data/val.txt'

os.makedirs(os.path.dirname(train_txt_path), exist_ok=True)
os.makedirs(os.path.dirname(val_txt_path), exist_ok=True)

supported_formats = ('.jpg', '.jpeg', '.png', '.bmp')
image_paths = [os.path.abspath(os.path.join(dp, f)) 
               for dp, dn, filenames in os.walk(image_dir) 
               for f in filenames if f.endswith(supported_formats)]

random.shuffle(image_paths)
split_index = int(train_ratio * len(image_paths))

train_image_paths = image_paths[:split_index]
val_image_paths = image_paths[split_index:]

with open(train_txt_path, 'w') as file:
    for image_path in train_image_paths:
        file.write(f"{image_path}\n")

with open(val_txt_path, 'w') as file:
    for image_path in val_image_paths:
        file.write(f"{image_path}\n")

print(f"Training images have been written to {train_txt_path}, total: {len(train_image_paths)}")
print(f"Validation images have been written to {val_txt_path}, total: {len(val_image_paths)}")

3、相关配置文件

数据集配置文件

创建包含数据集配置信息的 .yaml 文件，它提供了结构化的方法来存储训练集和验证集的划分文件路径，以及分类信息等其他与数据集相关的配置。以下是 mydata.yaml 文件的一个示例内容，其中包含 train.txt 和 val.txt 文件路径的条目，以及数据集的分类信息：

train: /home/detect/datasets/mydata/data/train.txt
val: /home/detect/datasets/mydata/data/val.txt

nc: 2
names: ['cat', 'dog']

YOLO 模型配置文件

在 ultralytics/cfg/models/v8/ 目录中，您可以找到各种 YOLOv8 模型的配置文件。您可以选择合适的模型配置文件，并根据自己的需求进行调整。

在本教程中，我们将遵循默认配置，因此不需要对其做任何更改。我们将直接使用这些预设参数，以便您能快速地上手并运行模型。

4、YOLO 模型训练

下载预训练权重

如果您的训练服务器能够顺畅访问 GitHub，在训练过程中，系统将会自动下载并载入相应的预训练模型权重进行初始化。

如若不然，则需要手动前往 Ultralytics Assets 下载您所需模型的预训练权重文件，并传输到训练服务器指定目录下。

开始模型训练

使用我们自己的数据集进行 YOLOv8n 模型训练，持续 1000 个 epoch。训练设备可以通过设备参数进行指定，默认将自动选择，若GPU不可用，则使用 CPU 作为训练设备，否则将使用第一个可用的 GPU。

=== “Python”

from ultralytics import YOLO

model = YOLO("yolov8n.yaml")  # 从 YAML 构建新模型
model = YOLO("yolov8n.pt")  # 加载预训练模型权重
model = YOLO("yolov8n.yaml").load("yolov8n.pt")  # YAML构建并加载权重

# 模型训练
results = model.train(data="mydata.yaml", epochs=1000, imgsz=640, batch=256)
# 使用多 GPU 进行模型训练
results = model.train(data="mydata.yaml", epochs=1000, imgsz=640, batch=256, device=[0, 1])

=== “CLI”

# 从 YAML 构建新模型，并开始模型训练
yolo detect train data=mydata.yaml model=yolov8n.yaml epochs=1000 imgsz=640 batch=256

# 加载预训练模型权重，并开始模型训练
yolo detect train data=mydata.yaml model=yolov8n.pt epochs=1000 imgsz=640 batch=256

# YAML构建并加载权重，同时开始模型训练
yolo detect train data=mydata.yaml model=yolov8n.yaml pretrained=yolov8n.pt epochs=1000 imgsz=640 batch=256

# YAML构建并加载权重，同时使用多 GPU 开始模型训练
yolo detect train data=mydata.yaml model=yolov8n.yaml pretrained=yolov8n.pt epochs=1000 imgsz=640 batch=256 device=0,1

YOLOv8 迎来了对 Apple M1 和 M2 芯片的优化支持，允许用户通过 Metal Performance Shaders (MPS) 在这些平台上进行高效的训练。此外，它还支持从中断点恢复训练，并提供了改进的日志记录功能。

有关更多的功能选项和训练参数的详细配置，您可以参阅【官方训练文档】以获取全面的指南和说明

5、YOLO 模型推理

Ultralytics YOLOv8 提供了具有强大功能的预测模式，该模式专为在各种数据源上实现高性能实时推理而设计。

在推理过程中，Ultralytics YOLO 模型会返回 Python Results 对象的列表，或者在传递 stream=True 参数时，返回内存高效的 Python Results 对象生成器。官方推理示例代码如下：

from ultralytics import YOLO

# Load a pretrained YOLOv8n model
model = YOLO("yolov8n.pt")

''' Run batched inference on a list of images'''
# return a list of Results objects
results = model(["im1.jpg", "im2.jpg"])
# # return a generator of Results objects
# results = model(["im1.jpg", "im2.jpg"], stream=True)

for result in results:
    boxes = result.boxes          # Boxes object for bounding box outputs
    masks = result.masks          # Masks object for segmentation masks outputs
    keypoints = result.keypoints  # Keypoints object for pose outputs
    probs = result.probs          # Probs object for classification outputs
    obb = result.obb              # Oriented boxes object for OBB outputs
    result.show()
    result.save(filename="result.jpg")

此外，您还可以详细处理得到的 bounding box 输出信息。以下是简单的示例代码：

detect_output = {"boxes": [], "scores": [], "classes": [], "names": [], "masks": []}

for result in results:
    detect_output["boxes"].append(result.boxes.xyxy.cpu())
    detect_output["scores"].append(result.boxes.conf.cpu())

    names = self.model.names
    classes = result.boxes.cls.cpu().numpy().astype(int)
    class_names = [names[c] for c in classes]
    detect_output["classes"].extend(classes)
    detect_output["names"].extend(class_names)

    if result.masks is not None:
        detect_output["masks"].extend(result.masks.numpy())

有关更多关于模型推理输入支持和结果对象的详细参数信息，请参阅【官方预测文档】以获取全面的指南和说明。

YOLOv8 深度学习 计算机视觉 CV Deep Learning