需求内容

之前写了一篇【目标检测】YOLOv5跑通VOC2007数据集里面写了个脚本是将xml标注格式转换称Yolo格式，同时读取数据集划分。
在训练自己的数据时，我发现没有现成的数据集划分文件，于是就写了这个小脚本来读取文件名，并划分成train/val/test三个文件。

需求实现

完整脚本：

import os
import random

img_path = './image/'
img_list = os.listdir(img_path)
train_ratio = 0.8  # 训练集比例
val_ratio = 0.1  # 验证集比例
shuffle = True  # 是否随机划分
def data_split(full_list, train_ratio, val_ratio, shuffle=True):
    n_total = len(full_list)
    train_set_num = int(n_total * train_ratio)
    val_set_num = int(n_total * val_ratio)
    if shuffle:
        random.shuffle(full_list)
    train_set = full_list[:train_set_num]
    val_set = full_list[train_set_num:(train_set_num + val_set_num)]
    test_set = full_list[(train_set_num + val_set_num):]
    return train_set, val_set, test_set
if __name__ == '__main__':
    train_set, val_set, test_set = data_split(img_list, train_ratio, val_ratio, shuffle=True)

    with open('train.txt', 'w') as f:
        for img_name in train_set:
            f.write(img_name.split('.png')[0] + '\n')

    with open('val.txt', 'w') as f:
        for img_name in val_set:
            f.write(img_name.split('.png')[0] + '\n')

    with open('test.txt', 'w') as f:
        for img_name in test_set:
            f.write(img_name.split('.png')[0] + '\n')

  
 

  
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
 

读取图片数据：

划分后的数据：

文章来源: zstar.blog.csdn.net，作者：zstar-_，版权归原作者所有，如需转载，请联系作者。

原文链接：zstar.blog.csdn.net/article/details/125831710