训练yolo8算法,记录下过程,避免每次训练都忘记
1.xml格式转txt

import os
import xml.etree.ElementTree as ET

def convert(size, box):
    """
    将边界框从Pascal VOC格式转换为YOLO格式。
    :param size: 图片尺寸 (width, height)
    :param box: Pascal VOC格式的边界框 (xmin, ymin, xmax, ymax)
    :return: YOLO格式的中心点坐标和宽高 (center_x, center_y, width, height)
    """
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[2]) / 2.0
    y = (box[1] + box[3]) / 2.0
    w = box[2] - box[0]
    h = box[3] - box[1]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h

def convert_annotation(xml_file_path, classes, output_dir):
    """
    解析XML文件并生成YOLO格式的标签。
    :param xml_file_path: XML文件路径
    :param classes: 类别列表
    :param output_dir: 输出目录
    """
    tree = ET.parse(xml_file_path)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    out_file_path = os.path.join(output_dir, os.path.basename(xml_file_path).replace('.xml', '.txt'))
    out_file = open(out_file_path, 'w')

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('xmax').text), float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

if __name__ == "__main__":
    xml_labels_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/xmls'  # XML标签所在目录
    classes_file = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/classes.txt'  # 类别文件路径
    output_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/test'  # 输出YOLO标签的目录

    # 加载类别
    with open(classes_file, 'r') as f:
        classes = [line.strip() for line in f.readlines()]

    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历所有XML文件并转换
    for xml_file in os.listdir(xml_labels_dir):
        if not xml_file.endswith('.xml'):
            continue
        xml_file_path = os.path.join(xml_labels_dir, xml_file)
        convert_annotation(xml_file_path, classes, output_dir)

    print("转换完成")

2.训练集、测试集转换代码

import os
import random
import shutil

def split_yolo_dataset(image_dir, label_dir, output_dir, train_ratio=0.8, val_ratio=0.2, test_ratio=0):
    """
    Splits the YOLO-formatted dataset into training, validation, and testing sets based on the provided ratios.
    """
    if not (train_ratio + val_ratio + test_ratio) == 1.0:
        raise ValueError("The sum of train_ratio, val_ratio, and test_ratio must be 1.0")

    # List all images without extensions
    images = [os.path.splitext(f)[0] for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]
    total_images = len(images)
    print(f"Total images found: {total_images}")

    # Shuffle images to ensure randomness
    random.shuffle(images)

    # Calculate sizes of each split using round to avoid rounding issues
    train_count = round(total_images * train_ratio)
    val_count = round(total_images * val_ratio)
    test_count = total_images - train_count - val_count

    # Adjust counts if necessary to ensure no images are left out
    if test_count < 0:
        val_count += test_count
        test_count = 0
    elif train_count + val_count + test_count != total_images:
        # In case of rounding error, adjust the count of the last set
        diff = total_images - (train_count + val_count + test_count)
        if test_count > 0:
            test_count += diff
        elif val_count > 0:
            val_count += diff
        else:
            train_count += diff

    # Split the dataset
    splits = {
        'train': images[:train_count],
        'val': images[train_count:train_count + val_count],
        'test': images[train_count + val_count:]
    }

    # Create output directories and copy/move files accordingly
    for split_name, split_images in splits.items():
        if not split_images:  # 如果该集合没有图像,则跳过创建目录
            continue
        split_image_dir = os.path.join(output_dir, split_name, 'images')
        split_label_dir = os.path.join(output_dir, split_name, 'labels')
        os.makedirs(split_image_dir, exist_ok=True)
        os.makedirs(split_label_dir, exist_ok=True)

        for img_base in split_images:
            # Find the correct extension for the image file
            img_ext = next(
                (ext for ext in ['.jpg', '.jpeg', '.png'] if os.path.exists(os.path.join(image_dir, img_base + ext))),
                '')
            if not img_ext:
                print(f"Warning: Could not find an image for base name {img_base}")
                continue

            src_img_path = os.path.join(image_dir, img_base + img_ext)
            dst_img_path = os.path.join(split_image_dir, img_base + img_ext)
            shutil.copy(src_img_path, dst_img_path)

            src_label_path = os.path.join(label_dir, img_base + '.txt')
            if os.path.exists(src_label_path):
                dst_label_path = os.path.join(split_label_dir, img_base + '.txt')
                shutil.copy(src_label_path, dst_label_path)
            else:
                print(f"Warning: No corresponding label file found for {img_base}")

if __name__ == "__main__":
    image_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/image'  # 替换为实际图像所在目录
    label_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/test'  # 替换为实际YOLO标签所在目录
    output_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/3.6-now'  # 替换为你希望存放分割后数据集的目录

    split_yolo_dataset(image_dir, label_dir, output_dir, train_ratio=0.8, val_ratio=0.2, test_ratio=0)

3.训练脚本

import warnings

warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':

    model = YOLO('/home/columbus/hlz/3.6/yolov8/ultralytics-main/ultralytics/cfg/models/v8/yolov8.yaml')  # 指定YOLO模型对象,并加载指定配置文件中的模型配置
    model.load('/home/columbus/hlz/3.6/yolov8/ultralytics-main/yolov8s.pt')      #加载预训练的权重文件'yolov8s.pt',加速训练并提升模型性能
    model.train(data='/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/date.yaml',  # 指定训练数据集的配置文件路径,这个.yaml文件包含了数据集的路径和类别信息
                epochs=150,  # 设置训练的总轮数为200轮
                batch=32,  # 设置每个训练批次的大小为16,即每次更新模型时使用16张图片
                #close_mosaic=10,  # 设置在训练结束前多少轮关闭 Mosaic 数据增强,10 表示在训练的最后 10 轮中关闭 Mosaic
                #workers=8,  # 设置用于数据加载的线程数为8,更多线程可以加快数据加载速度
                #patience=50,  # 在训练时,如果经过50轮性能没有提升,则停止训练(早停机制)
                #device='0',  # 指定使用的设备,'0'表示使用第一块GPU进行训练
                optimizer='SGD',  # 设置优化器为SGD(随机梯度下降),用于模型参数更新
                )

更多推荐