【目标检测】小脚本：数据集划分

zstar

发布于 2022-09-08 11:35:22

5750

发布于 2022-09-08 11:35:22

文章被收录于专栏：往期博文

需求内容

之前写了一篇【目标检测】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')

读取图片数据：