yolo数据集划分脚本代码

数据集标注完后，是不能直接用于训练的，需要对数据集进行一定比例划分为训练集和验证集，可以手动划分，但是这样比较麻烦，于是就有了代码划分数据集。

在如下的代码中，

• 把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("提取完毕")
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.

提取完毕后，它会在图片所在文件夹的同级目录，生成dataset文件夹

里面有划分好的训练集和验证集（原来的图片和标签会移动到这里）