从零到一训练YOLOv8模型

---

前言

YOLOv8 基于深度学习和计算机视觉领域的尖端技术，在速度和准确性方面具有无与伦比的性能。其流线型设计使其适用于各种应用，并可轻松适应从边缘设备到云 API 等不同硬件平台。

---

一、数据集准备

数据集文件存放规则：

以下是随机划分数据集的代码：

import os
import shutil
import random

# Directories
labels_dir = r"E:\File\VOC2020\output"
images_dir = r"E:\File\VOC2020\images"
total_dir = r"E:\File\VOC2020\fire_detect"

# Output directories
train_images_dir = os.path.join(total_dir, 'train', 'images')
train_labels_dir = os.path.join(total_dir, 'train', 'labels')
valid_images_dir = os.path.join(total_dir, 'valid', 'images')
valid_labels_dir = os.path.join(total_dir, 'valid', 'labels')
test_images_dir = os.path.join(total_dir, 'test', 'images')
test_labels_dir = os.path.join(total_dir, 'test', 'labels')

# Create output directories if they don't exist
os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(train_labels_dir, exist_ok=True)
os.makedirs(valid_images_dir, exist_ok=True)
os.makedirs(valid_labels_dir, exist_ok=True)
os.makedirs(test_images_dir, exist_ok=True)
os.makedirs(test_labels_dir, exist_ok=True)

# Get list of all image paths
image_paths = os.listdir(images_dir)
random.shuffle(image_paths)

# Split data into train, valid, and test sets
num_images = len(image_paths)
train_split = int(0.7 * num_images)
valid_split = int(0.2 * num_images)

train_image_paths = image_paths[:train_split]
valid_image_paths = image_paths[train_split:train_split + valid_split]
test_image_paths = image_paths[train_split + valid_split:]

# Function to copy images and labels


def copy_files(image_paths, dest_images_dir, dest_labels_dir):
    for image_path in image_paths:
        image_name = os.path.splitext(image_path)[0]
        xml_path = os.path.join(labels_dir, image_name + '.txt')
        if not os.path.exists(xml_path):
            continue

        # Copy image
        shutil.copy(os.path.join(images_dir, image_path), dest_images_dir)

        # Copy label
        shutil.copy(xml_path, dest_labels_dir)


# Copy files to respective directories
copy_files(train_image_paths, train_images_dir, train_labels_dir)
copy_files(valid_image_paths, valid_images_dir, valid_labels_dir)
copy_files(test_image_paths, test_images_dir, test_labels_dir)

print("Done")

把图片和标签区分开

def split_image_label(dir_path, image_dir, label_dir):
    # 获取目录中的所有文件名
    file_names = os.listdir(dir_path)

    for file_name in file_names:
        # 获取完整文件路径
        file_path = os.path.join(dir_path, file_name)

        # 获取文件扩展名
        _, ext = os.path.splitext(file_path)
        ext = ext.lower()

        if ext == '.xml' or ext==".txt":
            # 如果是XML文件，复制到label_dir目录
            dest_path = os.path.join(label_dir, file_name)
            shutil.copy(file_path, dest_path)
        elif ext == '.jpg' or ext == '.jpeg':
            # 如果是JPG文件，复制到image_dir目录
            dest_path = os.path.join(image_dir, file_name)
            shutil.copy(file_path, dest_path)
        else:
            # 其他类型文件，跳过处理
            pass

xml格式转换成txt格式：

import xml.etree.ElementTree as ET
import os
import shutil

labels = {"fire": 0}


def convert_to_yolo_format(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)

    yolo_annotations = []

    for obj in root.iter('object'):
        # Assuming class names are consistent and mapped to IDs
        class_id = obj.find('name').text
        class_id = labels.get(class_id)
        if class_id is None:
            continue  # Skip unknown classes

        bndbox = obj.find('bndbox')
        xmin = int(bndbox.find('xmin').text)
        ymin = int(bndbox.find('ymin').text)
        xmax = int(bndbox.find('xmax').text)
        ymax = int(bndbox.find('ymax').text)

        x_center = (xmin + xmax) / 2.0 / width
        y_center = (ymin + ymax) / 2.0 / height
        bbox_width = (xmax - xmin) / width
        bbox_height = (ymax - ymin) / height

        yolo_annotations.append(
            f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")

    return yolo_annotations


def create_yolo_file(yolo_annotations, output_file):
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    with open(output_file, 'a+') as f:
        for line in yolo_annotations:
            f.write(line + '\n')




images_dir = r"E:\File\VOC2020\images"
labels_dir = r"E:\File\VOC2020\labels"
output_dir = r"E:\File\VOC2020\output"
image_paths = os.listdir(images_dir)
for image_path in image_paths:
    image_name = os.path.splitext(image_path)[0]
    xml_path = os.path.join(labels_dir, image_name+'.xml')
    if not os.path.exists(xml_path):
        continue
    yolo_annotations = convert_to_yolo_format(xml_path)
    output_file = os.path.join(output_dir, image_name + '.txt')
    create_yolo_file(yolo_annotations, output_file)
print("Done")

二、训练模型

1.写好yaml文件

train: /root/product/fire_detect/train  # train images (relative to 'path') 128 images
val: /root/product/fire_detect/valid # val images (relative to 'path') 128 images
test: /root/product/fire_detect/test # test images (optional)

# Classes
names:
  0: fire

2.开始训练

from ultralytics import YOLO
import os

if __name__ == "__main__":
    # Load a model
    data_yaml = r"/root/product/data.yaml"
    model = YOLO("yolov8n.yaml") # build a new model from YAML
    model = YOLO("yolov8n.yaml").load("/root/product/yolov8n.pt") # build from YAML and transfer weights
    # model = YOLO(pre_model, task='detect')  # load a pretrained model (recommended for training)

    # Train the model
    results = model.train(data=data_yaml, epochs=200, imgsz=640, batch=16, workers=8, device=0,
    cos_lr=True, close_mosaic=200, warmup_epochs=10)

/root/product/yolov8n.pt是预训练模型需要自己去官网下

3.模型推理

from ultralytics import YOLO
import cv2
# Load a model
model = YOLO(r"/root/product/runs/detect/train3/weights/best.pt") # pretrained YOLOv8n model


image=cv2.imread("/root/product/1.jpg")
# Run batched inference on a list of images
results = model(image) # return a list of Results objects

# Process results list
for result in results:
    result.save(filename="result.jpg") # save to disk