简单的数据集划分python代码

目录

前言

内容简介

第一步：遍历标签文件和图片文件

第二步：统计新文件夹中文件名并存入对应列表

第三步：按照上一步生成的列表将文件存入对应文件

完整代码展示

---

前言

使用python代码实现简单的数据集分割，本文示例为jpg图片和labelme标注的json文件数据分割

内容简介

在我们对数据集完成标注后，如用labelme对数据集进行实例分割标注，能得到一些json格式标注文件，这些文件名和图片名称一一对应，在我们进行实例分割网络训练时，需要常需要将数据划分为训练集(train)、验证集(val)、测试集(testing)三个文件夹。以下进行数据集划分的简单代码，文中仅涉及train和val文件夹划分，testing同理，仅供参考。

第一步：遍历标签文件和图片文件

import os
import shutile
import random

random.seed(0) #随机数种子，确保每次运行结果一致
#定义分割函数，train_rate为训练集比例，val_rate为训练集比例
def split_data(file_path,label_path, new_file_path, train_rate, val_rate):

    each_class_image = []
    each_class_label = []
    for image in os.listdir(image_path): 
    each_class_image>append(image)  #遍历图片文件夹，记录图片文件名
    for label in os.listdir(label_path):
    each_class_label.append(label)  #遍历标签文件夹，记录标签文件名

第二步：统计新文件夹中文件名并存入对应列表

    data=list(zip(each_class_image,each_class_label))#图片和标签排序，绑定
        total = len(each_class_image) #记录图片（标签）文件个数
        random.shuffle(data)       #打乱文件顺序，但是同名标签和图片仍绑定为元组    
     #按顺序取出图片和标签文件名，存入两个列表
        each_class_image,each_class_label=zip(*data) 
     #统计各文件夹中将存入数据的文件名
        train_images = each_class_image[0:int(train_rate * total)]
        val_images = each_class_image[int(train_rate * total):int((train_rate + val_rate) *                     
    total)]
        train_labels = each_class_label[0:int(train_rate * total)]
        val_labels = each_class_label[int(train_rate * total):int((train_rate + val_rate) *     
    total)]

第三步：按照上一步生成的列表将文件存入对应文件

    for image in train_images:
        print(image)
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'train' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1) #创建文件夹
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)#从原文件夹复制文件到新文件夹

    for label in train_labels:
        print(label)
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'train' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)

    for image in val_images:
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'val' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in val_labels:
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'val' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)

完整代码展示

import os
import shutil
import random

random.seed(0)


def split_data(file_path,label_path, new_file_path, train_rate, val_rate):
    each_class_image = []
    each_class_label = []
    for image in os.listdir(file_path):
        each_class_image.append(image)
    for label in os.listdir(label_path):
        each_class_label.append(label)
    data=list(zip(each_class_image,each_class_label))
    total = len(each_class_image)
    random.shuffle(data)
    each_class_image,each_class_label=zip(*data) #与zip相反，可理解为解压
    train_images = each_class_image[0:int(train_rate * total)]
    val_images = each_class_image[int(train_rate * total):int((train_rate + val_rate) * total)]
    train_labels = each_class_label[0:int(train_rate * total)]
    val_labels = each_class_label[int(train_rate * total):int((train_rate + val_rate) * total)]

    for image in train_images:
        print(image)
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'train' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1) #创建目录
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)#从原文件夹复制文件到新文件夹

    for label in train_labels:
        print(label)
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'train' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)

    for image in val_images:
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'val' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in val_labels:
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'val' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)
if __name__ == '__main__':
    file_path = r'原来图片的地址'
    label_path = r'json标签文件的地址'
    new_file_path = r'生成新文件地址'
    split_data(file_path,xml_path, new_file_path, train_rate=0.7, val_rate=0.3) #训练验证集比例