在训练YOLO(You Only Look Once)目标检测模型时,自定义数据集的划分是一个重要步骤。数据集的划分主要包括训练集、验证集和测试集的划分。这样做有助于模型在训练过程中进行有效的学习和评估。
在yolov8项目文件夹下创建存储原始数据集的目录,比如:origin_dataset(可以是其它名字),在这个目录下创建labels和images这两个文件夹。
在yolov8项目文件夹下,创建split.py文件,用于存储划分数据的代码。
这里划分比例:训练集:验证集:测试集 = 7 : 1 :2。
split.py
import os
import random
import shutil
from pathlib import Path
#这里划分比例:训练集:验证集:测试集 = 7 : 1 :2。
def split_dataset(dataset_dir, output_dir, train_ratio=0.7, test_ratio=0.2, val_ratio=0.1):
"""
自定义数据集划分函数,按照 YOLO 数据集格式划分数据集
:param dataset_dir: 数据集的根目录(包含 images 和 labels 子文件夹)
:param output_dir: 划分后的数据集保存目录
:param train_ratio: 训练集占比
:param test_ratio: 测试集占比
:param val_ratio: 验证集占比
"""
# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)
# 创建子目录结构
train_images_dir = os.path.join(output_dir, "train", "images")
train_labels_dir = os.path.join(output_dir, "train", "labels")
val_images_dir = os.path.join(output_dir, "valid", "images")
val_labels_dir = os.path.join(output_dir, "valid", "labels")
test_images_dir = os.path.join(output_dir, "test", "images")
test_labels_dir = os.path.join(output_dir, "test", "labels")
os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(train_labels_dir, exist_ok=True)
os.makedirs(val_images_dir, exist_ok=True)
os.makedirs(val_labels_dir, exist_ok=True)
os.makedirs(test_images_dir, exist_ok=True)
os.makedirs(test_labels_dir, exist_ok=True)
# 获取所有图像文件路径
image_files = []
for ext in ['.jpg', '.png', '.jpeg']: # 支持的图像格式
image_files.extend(list(Path(dataset_dir).rglob(f'images/*{ext}')))
# 打乱文件列表
random.shuffle(image_files)
# 计算划分数量
total_files = len(image_files)
train_size = int(total_files * train_ratio)
test_size = int(total_files * test_ratio)
val_size = int(total_files * val_ratio)
# 划分数据集
train_files = image_files[:train_size]
val_files = random.sample(train_files, val_size) # 从训练集中随机抽取验证集
test_files = image_files[train_size:train_size + test_size]
# 复制文件到对应目录
def copy_files(file_list, images_dir, labels_dir):
for file in file_list:
#print("file=>",file)
# 复制图像文件
shutil.copy(file, images_dir)
# 获取对应的标签文件路径
image_name = os.path.basename(file)
image_name_without_ext = os.path.splitext(image_name)[0]
label_file = os.path.join(dataset_dir, "labels", f"{image_name_without_ext}.txt")
if os.path.exists(label_file):
shutil.copy(label_file, labels_dir)
else:
print(f"警告:标签文件 {label_file} 不存在,跳过复制。")
# 复制训练集文件
#print("train_images_dir:",train_images_dir,"train_labels_dir:",train_labels_dir)
copy_files(train_files, train_images_dir, train_labels_dir)
print(f"训练集: {len(train_files)} 张图片")
# 复制验证集文件
copy_files(val_files, val_images_dir, val_labels_dir)
print(f"验证集: {len(val_files)} 张图片")
# 复制测试集文件
copy_files(test_files, test_images_dir, test_labels_dir)
print(f"测试集: {len(test_files)} 张图片")
print(f"数据集划分完成!")
# 示例调用
dataset_dir = "./origin_dataset" # 数据集的根目录(包含 images 和 labels 子文件夹)
output_dir = "./output/dataset" # 划分后的数据集保存目录
split_dataset(dataset_dir, output_dir)
运行效果:
训练集: 2488 张图片
验证集: 355 张图片
测试集: 711 张图片
数据集划分完成!
运行后,划分后的数据集自动保存在了项目的output目录下,如下图所示:
