数据集标注完后,是不能直接用于训练的,需要对数据集进行一定比例划分为训练集和验证集,可以手动划分,但是这样比较麻烦,于是就有了代码划分数据集。
在如下的代码中,
- 把src_img_dir 的值写入图片文件夹的路径
- 把src_label_dir 的值,写入txt标签文件夹的路径
然后运行代码即可划分,生成dataset文件夹
import os
import shutil
import random
# 源目录
src_img_dir = r'G:\down\score\score\images'
src_label_dir = r'G:\down\score\score\labels'
# 创建dataset目录在与原图片文件夹同级目录
dataset_dir = os.path.join(os.path.dirname(src_img_dir), 'dataset')
train_dir = os.path.join(dataset_dir, 'train')
val_dir = os.path.join(dataset_dir, 'val')
os.makedirs(os.path.join(train_dir, 'images'), exist_ok=True)
os.makedirs(os.path.join(train_dir, 'labels'), exist_ok=True)
os.makedirs(os.path.join(val_dir, 'images'), exist_ok=True)
os.makedirs(os.path.join(val_dir, 'labels'), exist_ok=True)
# 获取源目录下所有文件名
img_filenames = os.listdir(src_img_dir)
label_filenames = os.listdir(src_label_dir)
# 乱序文件列表
random.shuffle(img_filenames)
# 计算训练集和验证集数量
num_images = len(img_filenames)
num_train = int(0.8 * num_images)
num_val = num_images - num_train
# 分割数据
train_images = img_filenames[:num_train]
val_images = img_filenames[num_train:]
for filename in train_images:
# 检查文件是否为.jpg或.png图片
if filename.lower().endswith(('.jpg', '.png','.tif','jpeg','.webp')):
src_img_path = os.path.join(src_img_dir, filename)
dst_img_path = os.path.join(train_dir, 'images', filename)
base_name, ext = os.path.splitext(filename)
label_filename = base_name + '.txt'
src_label_path = os.path.join(src_label_dir, label_filename)
dst_label_path = os.path.join(train_dir, 'labels', label_filename)
# 检查源文件是否存在
if os.path.exists(src_img_path) and os.path.exists(src_label_path):
shutil.move(src_img_path, dst_img_path)
shutil.move(src_label_path, dst_label_path)
for filename in val_images:
# 检查文件是否为.jpg或.png图片
if filename.lower().endswith(('.jpg', '.png','.tif','jpeg','.webp')):
src_img_path = os.path.join(src_img_dir, filename)
dst_img_path = os.path.join(val_dir, 'images', filename)
base_name, ext = os.path.splitext(filename)
label_filename = base_name + '.txt'
src_label_path = os.path.join(src_label_dir, label_filename)
dst_label_path = os.path.join(val_dir, 'labels', label_filename)
# 检查源文件是否存在
if os.path.exists(src_img_path) and os.path.exists(src_label_path):
shutil.move(src_img_path, dst_img_path)
shutil.move(src_label_path, dst_label_path)
print("提取完毕")
提取完毕后,它会在图片所在文件夹的同级目录,生成dataset文件夹
里面有划分好的训练集和验证集(原来的图片和标签会移动到这里)