介绍
- 代码按照
8:2
的比例划分数据集为训练集和验证集,具体比例可以修改split_ratio
参数自行调整。 split_ratio
:数据集划分比例,默认为 0.8,表示将80%
的数据用于训练集,20%
用于验证集。- 划分的数据包括
jpg图像文件
和txt文件
,用于YOLOv8模型训练。
Python实现
import os, random, shutil
"""
按照8:2的比例划分数据集为训练集和验证集
"""
def split_dataset(img_path, txt_path, tar_imgs_path, tar_labels_path, split_ratio=0.8):
"""
将数据集划分为训练集和验证集,并保存到相应的文件夹中。
Parameters:
- split_ratio: 数据集划分比例,默认为 0.8,表示将 80% 的数据用于训练集,20% 用于验证集。
"""
os.makedirs(os.path.join(tar_imgs_path, "train"), exist_ok=True)
os.makedirs(os.path.join(tar_imgs_path, "val"), exist_ok=True)
os.makedirs(os.path.join(tar_labels_path, "train"), exist_ok=True)
os.makedirs(os.path.join(tar_labels_path, "val"), exist_ok=True)
# 获取数据集中所有文件的列表
img_file_list = os.listdir(img_path)
random.shuffle(img_file_list)
# 根据划分比例计算训练集和验证集的边界索引
split_index = int(len(img_file_list) * split_ratio)
print("split_index:", split_index)
train_img = img_file_list[:split_index]
val_img = img_file_list[split_index:]
# 将训练集数据移动到相应文件夹
for file in train_img:
if file.endswith('.jpg'):
img_src = os.path.join(img_path, file)
label_src = os.path.join(txt_path, file[:-4] + '.txt')
shutil.copy(img_src, os.path.join(os.path.join(tar_imgs_path, "train", file)))
shutil.copy(label_src, os.path.join(os.path.join(tar_labels_path, "train", file[:-4] + '.txt')))
# 将验证集数据移动到相应文件夹
for file in val_img:
if file.endswith('.jpg'):
img_src = os.path.join(img_path, file)
label_src = os.path.join(txt_path, file[:-4] + '.txt')
shutil.copy(img_src, os.path.join(os.path.join(tar_imgs_path, "val", file)))
shutil.copy(label_src, os.path.join(os.path.join(tar_labels_path, "val", file[:-4] + '.txt')))
if __name__ == '__main__':
# 存放图片的文件夹和存放txt文件的总文件夹路径
img_path = r"E:\xxxx\images"
txt_path = r"E:\xxxx\labels"
# 目标图片的文件夹和存放txt文件的文件夹路径
tar_imgs_path = r'E:\fireData\images'
tar_labels_path = r'E:\fireData\labels'
# 调用函数划分数据集
split_dataset(img_path, txt_path, tar_imgs_path, tar_labels_path)
print("划分数据集成功!")
文件夹格式为:
- fireData
(自己存放数据的文件夹,自己命名)
- images
- train
- val
- labels
- train
- val
- images
划分好数据集之后,就可以将fireData压缩成zip包,上传到远程服务器解压缩,并开始训练步骤。
整理不易🚀🚀,关注和收藏后拿走📌📌欢迎留言🧐👋📣✨
快来关注我的公众号🔎AdaCoding 和 GitHub🔎 AdaCoding123