YOLOv8训练自己的数据集(记录)

最新推荐文章于 2024-08-19 11:28:26 发布

爱编码的小陈

最新推荐文章于 2024-08-19 11:28:26 发布

阅读量1.9k

点赞数 5

分类专栏：深度学习文章标签： YOLO

本文链接：https://blog.csdn.net/m0_70694811/article/details/136906507

版权

深度学习专栏收录该内容

10 篇文章 0 订阅

订阅专栏

一、准备前的文件夹目录介绍

bag-images文件夹：用来存放原始数据集所有的.jpg图片

xml文件夹:用来存放原始数据集打过标签的所有xml文件

txt文件夹:用来存放原始数据集，由xml格式转换为txt格式的所有文件

bag文件夹：是我们目标制作的数据集，用于后期跑实验

bag文件夹下有images和labels文件夹，每个文件夹下都有一个train和val文件夹
images文件夹：用来存放目标数据集的所有图片，分为train训练和val验证两部分

labels文件夹：用来存放目标数据集的所有对应的txt文件，分为train训练和val验证两部分

二、首先将自己标注完成的数据集进行格式转换

下面是将文件由xml格式转换为txt格式代码，建一个voc_label.py脚本，将下面代码复制进去，运行，生成（注意该文件路径地址）：

import xml.etree.ElementTree as ET
import os

classes = ["bag"]  # 类别,改成自己的类别名称

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))


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


def convert_annotation(image_id):
    in_file = open(r'E:\人包物\bag\xml/%s.xml' % (image_id), encoding='UTF-8')#E:\人包物\bag\xml    为xml文件地址

    out_file = open(r'E:\人包物\bag\txt/%s.txt' % (image_id), 'w')  # 生成txt格式文件  #'E:\人包物\bag\txt  为将要输出生成的txt文件地址
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        cls = obj.find('name').text
        # print(cls)
        if cls not in classes:
            continue
        cls_id = classes.index(cls)
        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))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

xml_path = os.path.join(CURRENT_DIR, r'E:\人包物\bag\xml/')  #E:\人包物\bag\xml    为xml文件地址

# xml list
img_xmls = os.listdir(xml_path)
for img_xml in img_xmls:
    label_name = img_xml.split('.')[0]
    print(label_name)
    convert_annotation(label_name)

三、将处理好的图片和txt文本进行train和val数据集划分

import os  # 用于处理文件路径、创建目录等操作
import random  # 用于生成随机数种子、打乱列表等操作
import shutil  # 用于生成随机数种子、打乱列表等操作
# 设置随机数种子
random.seed(123)
# 定义文件夹路径
image_dir = r'E:\人包物\images'  # 原始图像所在的子目录
label_dir = r'E:\人包物\labels'  # 原始标签所在的子目录
output_dir = r'E:\人包物\bag'  # 处理后的数据集输出目录
# 定义训练集、验证集和测试集比例
train_ratio = 0.8  # 训练集比例
valid_ratio = 0.1  # 验证集比例
test_ratio = 0.1  # 测试集比例
# 获取所有图像文件和标签文件的文件名（不包括文件扩展名）
image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]  # 提取所有图像文件的文件名列表
label_filenames = [os.path.splitext(f)[0] for f in os.listdir(label_dir)]  # 提取所有标签文件的文件名列表
# 随机打乱文件名列表
random.shuffle(image_filenames)  # 打乱图像文件的文件名列表
# 计算训练集、验证集和测试集的数量
total_count = len(image_filenames)  # 总文件数
train_count = int(total_count * train_ratio)  # 训练集文件数
valid_count = int(total_count * valid_ratio)  # 验证集文件数
test_count = total_count - train_count - valid_count  # 测试集文件数
# 定义输出文件夹路径
train_image_dir = os.path.join(output_dir, 'train', 'images')  # 训练集图像输出目录
train_label_dir = os.path.join(output_dir, 'train', 'labels')  # 训练集标签输出目录
valid_image_dir = os.path.join(output_dir, 'valid', 'images')  # 验证集图像输出目录
valid_label_dir = os.path.join(output_dir, 'valid', 'labels')  # 验证集标签输出目录
test_image_dir = os.path.join(output_dir, 'test', 'images')  # 测试集图像输出目录
test_label_dir = os.path.join(output_dir, 'test', 'labels')  # 测试集标签输出目录
# 创建输出文件夹
os.makedirs(train_image_dir, exist_ok=True)  # 创建训练集图像输出目录
os.makedirs(train_label_dir, exist_ok=True)  # 创建训练集标签输出目录
os.makedirs(valid_image_dir, exist_ok=True)  # 创建验证集图像输出目录
os.makedirs(valid_label_dir, exist_ok=True)  # 创建验证集标签输出目录
os.makedirs(test_image_dir, exist_ok=True)  # 创建测试集图像输出目录
os.makedirs(test_label_dir, exist_ok=True)  # 创建测试集标签输出目录
# 将图像和标签文件划分到不同的数据集中
for i, filename in enumerate(image_filenames):
    # 如果文件数量小于训练数据集大小，则将文件复制到训练数据集目录中
    if i < train_count:
        output_image_dir = train_image_dir
        output_label_dir = train_label_dir
    # 如果文件数量小于训练数据集大小+验证数据集大小，则将文件复制到验证数据集目录中
    elif i < train_count + valid_count:
        output_image_dir = valid_image_dir
        output_label_dir = valid_label_dir
    # 否则，将文件复制到测试数据集目录中
    else:
        output_image_dir = test_image_dir
        output_label_dir = test_label_dir
    # 复制图像文件
    src_image_path = os.path.join(image_dir, filename + '.jpg')  # 获取图像文件的源路径
    dst_image_path = os.path.join(output_image_dir, filename + '.jpg')  # 获取图像文件的目标路径
    shutil.copy(src_image_path, dst_image_path)  # 复制图像文件到目标路径
    # 复制标签文件
    src_label_path = os.path.join(label_dir, filename + '.txt')  # 获取标签文件的源路径
    dst_label_path = os.path.join(output_label_dir, filename + '.txt')  # 获取标签文件的目标路径
    shutil.copy(src_label_path, dst_label_path)  # 复制标签文件到目标路径

划分好之后就是下面这个样子：

将整个bag文件夹，复制移动到D:\cs\yolov8\ultralytics\datasets文件夹下，（一定要在此文件夹）

在D:\cs\yolov8\ultralytics\datasets文件夹下新建一个bag.yaml文件

# Ultralytics YOLO 🚀, AGPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: yolo train data=person.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── coco128  ← downloads here (7 MB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: D:/cs/yolov8/ultralytics/datasets/bag  # dataset root dir
train: images/train  # train images (relative to 'path') 128 images
val: images/val  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: bag

下载与训练模型

https://github.com/ultralytics/assets/releases

我们选择yolov8n.pt，权重下载好放在根目录下

接下来在根目录创建一个train.py脚本文件，将下面代码复制过去，进行训练

from ultralytics import YOLO

# 加载模型
# model = YOLO("yolov8n.yaml")  # 从头开始构建新模型
model = YOLO("yolov8n.pt")  # 加载预训练模型（推荐用于训练）

# Use the model
results = model.train(data=r"D:\cs\yolov8\ultralytics\datasets\bag.yaml", epochs=200, batch=16, )  # 训练模型