记录-YOLOv8训练自己的数据集

1.数据集制作

安装labelimg标注工具

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

启动labelimg

labelImg

调整模式，变成PascalVOC进行标注

?label快捷键

CTRL + S 保存标注
w 画框
a 上一张图片
d 下一张图片

voc标注方式，生成的是xml格式文件，需要转换为txt格式，转换代码如下：

import os
import xml.etree.ElementTree as ET

# XML 文件夹路径和输出 TXT 文件夹路径
xml_folder = 'data'
txt_folder = 'data'

# 类别名称和对应的 ID 映射关系
class_mapping = {
    'mouse': 0,
    # 添加更多类别...
}

# 遍历 XML 文件夹中的每个文件
for xml_file in os.listdir(xml_folder):
    if not xml_file.endswith('.xml'):
        continue

    # 解析 XML 文件
    tree = ET.parse(os.path.join(xml_folder, xml_file))
    root = tree.getroot()

    # 获取图像的宽度和高度
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)

    # 创建相应的 TXT 文件
    txt_file = os.path.splitext(xml_file)[0] + '.txt'
    with open(os.path.join(txt_folder, txt_file), 'w') as f:
        # 遍历 XML 文件中的每个目标
        for obj in root.findall('object'):
            # 获取目标的类别和边界框信息
            class_name = obj.find('name').text
            class_id = class_mapping.get(class_name)
            bbox = obj.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)

            # 将边界框坐标转换为 YOLO 格式的比例
            x_center = (xmin + xmax) / 2 / width
            y_center = (ymin + ymax) / 2 / height
            bbox_width = (xmax - xmin) / width
            bbox_height = (ymax - ymin) / height

            # 将目标信息写入 TXT 文件中
            f.write(f'{class_id} {x_center} {y_center} {bbox_width} {bbox_height}\n')

建立yolov8数据集，数据集文件布局如下：

--DataSet
	--train
		--images
		--labels
	--test
		--images
		--labels
	--valid
		--images
		--labels
	--data.yaml
--yolov8 源码位置

数据集按照如上布局拆分代码如下：

import os
import random
import shutil

# 设置随机数种子
random.seed(42)

# 数据集文件夹路径和输出文件夹路径
data_folder = 'data'
train_folder = 'dataset/train'
test_folder = 'dataset/test/'
val_folder = 'dataset/val/'

# 计算每个子集的大小
total_files = len(os.listdir(os.path.join(data_folder, 'images')))
train_size = int(total_files * 0.7)
test_size = int((total_files - train_size) * 0.5)
val_size = total_files - train_size - test_size

# 获取所有图像文件的文件名列表
image_files = os.listdir(os.path.join(data_folder, 'images'))
random.shuffle(image_files)

# 复制图像和标注文件到相应的子集文件夹中
for i, image_file in enumerate(image_files):
    base_file_name = os.path.splitext(image_file)[0]
    image_path = os.path.join(data_folder, 'images', image_file)
    label_path = os.path.join(data_folder, 'labels', base_file_name + '.txt')

    if i < train_size:
        print(image_path)
        print(os.path.join(train_folder, 'images'))
        shutil.copy(image_path, os.path.join(train_folder, 'images'))
        shutil.copy(label_path, os.path.join(train_folder, 'labels'))
    elif i < train_size + test_size:
        shutil.copy(image_path, os.path.join(test_folder, 'images'))
        shutil.copy(label_path, os.path.join(test_folder, 'labels'))
    else:
        shutil.copy(image_path, os.path.join(val_folder, 'images'))
        shutil.copy(label_path, os.path.join(val_folder, 'labels'))

2.yolov8训练

环境准备

下载项目代码

git clone https://github.com/ultralytics/ultralytics

配置文件

将数据集放至项目根目录下，并在数据集同目录下配置data.yaml

train: D:\datasets\train
val: D:\datasets\valid 
test: D:\datasets\test

# Classes
names:
  0: mouse

修改yolov8.yaml文件（主要修改类别数）

模型训练

下载预训练模型至项目根目录

开始训练

yolo task=detect mode=train model=yolov8.yaml pretrained=yolov8n.pt epochs=10 batch=1 data=datasets/data.yaml

参考文档：

实践出真知——YOLOv8训练自己的数据集

完整且详细的Yolov8复现+训练自己的数据集

使用Yolov8训练自己的数据集