在目标检测任务中,数据集的划分通常分为训练集和验证集,以便在训练模型时评估模型的性能。这个过程对于有效训练和评估目标检测模型非常重要。下面是划分目标检测数据集的一般步骤:``
-
数据集组织:
确保你的数据集中包含图像文件(通常为 ".jpg" 或 ".png")和相应的标签文件,用于描述图像中的目标位置和类别。标签文件可以是类似于 YOLO 格式的 ".txt" 文件,也可以是其他目标检测任务中使用的格式。 -
划分比例选择:
确定训练集和验证集的划分比例。通常,常见的比例为 80-90% 的数据用于训练,而剩下的 10-20% 用于验证。这个比例可以根据你的数据量和任务的特性进行调整。 -
数据集划分代码:
编写代码来将数据集划分为训练集和验证集。这通常涉及将图像文件和相应的标签文件分别移动到两个不同的文件夹中。 -
文件夹结构:
为了方便管理,创建一个包含两个子文件夹(例如 "images" 和 "labels")的文件夹,分别用于存储图像和标签文件。训练集和验证集各自都应该有这样的文件夹结构。 -
随机化和保持一一对应:
在划分数据集之前,通常会对数据集进行随机化处理,以确保训练集和验证集中的样本都是随机选择的。在划分过程中,要保持图像和标签的一一对应关系,以确保训练和验证的一致性。 -
数据集划分结果验证:
验证划分的结果是否符合预期。你可以检查训练集和验证集的文件数量,确保图像和标签仍然保持一一对应,并且确保文件夹结构正确。
!!! 注意以下代码中all_data应含有图片与对应标签(在同一个文件夹中,且名称相同仅拓展名不同,分为.jpg和.txt文件)
import os
import random
import shutil
def split_dataset(srcDir, trainDir, valDir, split_ratio=0.9):
"""
将数据集划分为训练集和验证集,并保存到相应的文件夹中。
Parameters:
- srcDir: 原始数据集文件夹路径,包含图像和标签文件。
- trainDir: 训练集文件夹路径,包含 'images' 和 'labels' 子文件夹。
- valDir: 验证集文件夹路径,包含 'images' 和 'labels' 子文件夹。
- split_ratio: 数据集划分比例,默认为 0.9,表示将 90% 的数据用于训练集,10% 用于验证集。
"""
os.makedirs(os.path.join(trainDir, 'images'), exist_ok=True)
os.makedirs(os.path.join(trainDir, 'labels'), exist_ok=True)
os.makedirs(os.path.join(valDir, 'images'), exist_ok=True)
os.makedirs(os.path.join(valDir, 'labels'), exist_ok=True)
# 获取数据集中所有文件的列表
file_list = os.listdir(srcDir)
random.shuffle(file_list)
# 根据划分比例计算训练集和验证集的边界索引
split_index = int(len(file_list) * split_ratio)
train_files = file_list[:split_index]
val_files = file_list[split_index:]
# 将训练集数据移动到相应文件夹
for file in train_files:
if file.endswith('.jpg'):
img_src = os.path.join(srcDir, file)
label_src = os.path.join(srcDir, file[:-4] + '.txt')
shutil.move(img_src, os.path.join(trainDir, 'images', file))
shutil.move(label_src, os.path.join(trainDir, 'labels', file[:-4] + '.txt'))
# 将验证集数据移动到相应文件夹
for file in val_files:
if file.endswith('.jpg'):
img_src = os.path.join(srcDir, file)
label_src = os.path.join(srcDir, file[:-4] + '.txt')
shutil.move(img_src, os.path.join(valDir, 'images', file))
shutil.move(label_src, os.path.join(valDir, 'labels', file[:-4] + '.txt'))
if __name__ == '__main__':
# 输入文件夹路径
srcDir = r"C:\Users\86159\Desktop\hat\all_data"
trainDir = r'C:\Users\86159\Desktop\hat\train'
valDir = r'C:\Users\86159\Desktop\hat\val'
# 调用函数划分数据集
split_dataset(srcDir, trainDir, valDir)
评论区