YOLO系列模型训练数据准备：完整指南与脚本合集

最新推荐文章于 2025-03-04 13:52:42 发布

Charnlee_

最新推荐文章于 2025-03-04 13:52:42 发布

阅读量2.4k

点赞数 28

分类专栏： YOLO系列文章标签： YOLO 目标跟踪人工智能目标检测

本文链接：https://blog.csdn.net/qq_42987058/article/details/144946798

版权

YOLO系列专栏收录该内容

2 篇文章

订阅专栏

YOLO（You Only Look Once）是当前目标检测领域的主流算法之一。为了高效地训练YOLO模型，数据准备至关重要。本文将为您详细讲解从数据收集到数据标注、格式转换、数据集划分、配置文件编写及数据完整性验证的全流程，并提供相关脚本实现。

1. 数据收集

在数据收集阶段，需要明确目标检测任务，并收集相关图片数据。

明确任务目标

确定需要检测的目标种类及其标签（如车辆、人、动物等）。

数据来源

自己拍摄图片：适用于特定场景的数据需求。
公开数据集：如 COCO、Pascal VOC 和 OpenImages。
网络爬取：可使用爬虫工具获取数据。
生成图片：利用生成模型如 DALL·E 生成。
购买图片：如果有钱的情况下，可采购高质量数据。

2. 数据标注

标注是目标检测数据准备的重要环节。推荐以下工具：

标注工具

LabelImg：适用于 XML 格式（Pascal VOC），操作简便(我一般使用这个)。
LabelMe：支持 JSON 格式标注。
Roboflow：在线标注工具，支持多种格式。

标注步骤

绘制目标框：使用工具在图片中标注目标的边界框，并分配类别标签。
导出标注文件：保存为 XML 或其他支持的格式。

XML到YOLO格式的转换

将标注的 XML 文件转换为 YOLO 格式的 TXT 文件是训练必不可少的步骤。以下是转换脚本：

import os
import xml.etree.ElementTree as ET

# 定义类别
classes = ["Dog", "Cat", "Pig", "Fox", "Tiger", "Lion"]

# 输入和输出路径
xml_folder = "path/to/xml/folder"  # XML 文件所在目录
txt_folder = "path/to/txt/folder"  # 转换后的 TXT 文件保存目录

# 确保输出目录存在
os.makedirs(txt_folder, exist_ok=True)

# 解析 XML 文件并转换为 YOLO 格式
def convert_xml_to_txt(xml_file, txt_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    image_width = int(root.find("size/width").text)
    image_height = int(root.find("size/height").text)
    
    with open(txt_file, "w") as f:
        for obj in root.findall("object"):
            class_name = obj.find("name").text
            if class_name not in classes:
                continue
            class_id = classes.index(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) / image_width
            y_center = ((ymin + ymax) / 2) / image_height
            bbox_width = (xmax - xmin) / image_width
            bbox_height = (ymax - ymin) / image_height

            f.write(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}\n")

# 批量转换
for xml_file in os.listdir(xml_folder):
    if xml_file.endswith(".xml"):
        xml_path = os.path.join(xml_folder, xml_file)
        txt_path = os.path.join(txt_folder, os.path.splitext(xml_file)[0] + ".txt")
        convert_xml_to_txt(xml_path, txt_path)

print("XML 转换为 YOLO TXT 格式完成！")

3. 数据组织

YOLO标准数据格式

dataset/
├── images/
│   ├── train/
│   └── val/
├── labels/
│   ├── train/
│   └── val/
├── train.txt
├── val.txt
├── dataset.yaml

images/：存放图片。
labels/：存放与图片同名的标注文件（YOLO 格式）。
train.txt：存放训练集的图片路径。
val.txt：存放验证集的图片路径。

图片数据集划分

以下脚本可将图片随机划分为训练集和验证集：

import os
import shutil
import random

# 设置源文件夹和目标文件夹
source_folder = 'path/to/your/folder'  # 原始数据集文件夹
train_folder = 'path/to/your/folder'  # 训练集文件夹
val_folder = 'path/to/your/folder'  # 验证集文件夹

# 创建目标文件夹，如果不存在
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)

# 获取所有文件的列表
all_files = [f for f in os.listdir(source_folder) if f.endswith(('.jpg', '.jpeg', '.png'))]

# 排除以 'neg' 开头的文件
valid_files = [f for f in all_files if not f.startswith('neg')]

# 设置训练集和验证集的比例
train_ratio = 0.8
val_ratio = 0.2

# 随机打乱文件列表
random.shuffle(valid_files)

# 计算训练集和验证集的数量
train_count = int(len(valid_files) * train_ratio)
val_count = len(valid_files) - train_count

# 分割文件为训练集和验证集
train_files = valid_files[:train_count]
val_files = valid_files[train_count:]

# 将文件复制到目标文件夹
def copy_files(file_list, destination_folder):
    for file in file_list:
        source_path = os.path.join(source_folder, file)
        destination_path = os.path.join(destination_folder, file)
        shutil.move(source_path, destination_path)

# 复制训练集和验证集文件
copy_files(train_files, train_folder)
copy_files(val_files, val_folder)

print(f"训练集包含 {len(train_files)} 张图片，验证集包含 {len(val_files)} 张图片。")

根据图片名划分标签

根据图片名将labels拆分成同样的train和val文件,替换以下文件夹位置

import os
import shutil

# 定义源目录和目标目录
images_dir = 'path/to/your/folder' # 你拆分好的image文件夹位置
labels = 'path/to/your/folder' # 你转换好的txt标签文件夹位置
target_labels_dir = 'path/to/your/folder' # 输出为train和val文件夹格式的labels位置

# 定义 train 和 val 子目录
train_image_dir = os.path.join(images_dir, 'train')
val_image_dir = os.path.join(images_dir, 'val')

train_label_dir = os.path.join(target_labels_dir, 'train')
val_label_dir = os.path.join(target_labels_dir, 'val')

# 创建目标目录（如果不存在）
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(val_label_dir, exist_ok=True)

# 获取 images/train 和 images/val 中的文件名
train_images = os.listdir(train_image_dir)
val_images = os.listdir(val_image_dir)

# 复制对应的标注文件到 labels/train 和 labels/val
for image in train_images:
    # 假设图片和标注文件的名字相同
    label_name = image.split('.')[0] + '.txt'
    label_path = os.path.join(labels, label_name)
    if os.path.exists(label_path):
        shutil.copy(label_path, os.path.join(train_label_dir, label_name))

for image in val_images:
    # 假设图片和标注文件的名字相同
    label_name = image.split('.')[0] + '.txt'
    label_path = os.path.join(labels, label_name)
    if os.path.exists(label_path):
        shutil.copy(label_path, os.path.join(val_label_dir, label_name))

print("标注文件已按类别拆分。")

输出图片训练集和验证集的路径txt文件

将以下文件夹替换为你的image-train文件夹，即可输出images-train的路径txt文件，image-val同理

import os

def extract_image_paths(image_dir, output_file):
    """将图片文件路径提取并写入txt文件"""
    # 获取指定目录下的所有文件
    image_files = [f for f in os.listdir(image_dir) if os.path.isfile(os.path.join(image_dir, f))]
    
    # 将每个图片的完整路径写入txt文件
    with open(output_file, 'w') as file:
        for image in image_files:
            # 获取图片的完整路径
            image_path = os.path.join(image_dir, image)
            file.write(image_path + '\n')
    
    print(f"Total {len(image_files)} image paths written to {output_file}")

# 使用示例
image_dir = 'path/to/your/folder'  # 替换为你的图片文件夹路径
output_file = 'path/to/your/folder'  # 输出txt文件名称
extract_image_paths(image_dir, output_file)

4. 创建dataset.yaml配置文件

将train和val替换为前面转换好的image-train和image-val路径txt文件
nc为你的全部类别数

更改标签类别，注意类别的顺序别搞错了，这个很重要，因为牵扯到标签的索

train: path/to/dataset/images/train.txt
val: path/to/dataset/images/val.txt

nc: 3  # 类别数
names: ["Dog", "Cat", "Pig", "Fox", "Tiger", "Lion"]  # 类别名称

5. 数据验证与配置

验证类别数据

通过以下脚本检查每个类别的样本数量：

from collections import defaultdict

labels_path = 'path/to/labels'
classes = ["Dog", "Cat", "Pig", "Fox", "Tiger", "Lion"]
category_counts = defaultdict(int)

for txt_file in os.listdir(labels_path):
    if txt_file.endswith('.txt'):
        with open(os.path.join(labels_path, txt_file)) as f:
            for line in f:
                cls_id = int(line.split()[0])
                category_counts[classes[cls_id]] += 1

for cls, count in category_counts.items():
    print(f"{cls}: {count}")

6.开始训练

from ultralytics import YOLO

# 选择以下任意一种方式加载模型
model = YOLO("yolo11n.yaml")  # 从模型YMAL文件训练
model = YOLO("yolo11n.pt")  # 加载模型权重进行训练
model = YOLO("yolo11n.yaml").load("yolo11n.pt")  # 将预训练模型转移到yaml文件上进行训练

# 将以下dat换成你的yaml文件即可训练模型，YOLO代会根据你的yaml文件自动加载数据
results = model.train(data="dataset.yaml", epochs=100, imgsz=640)

通过以上流程，您可以高效准备YOLO训练数据，并快速开始模型训练。如需进一步优化或调整，欢迎留言交流！

！！！如果对你有帮助，给我点个赞吧😀