【数据集操作】划分数据集为VOC数据格式

uu分享的一个好链接：VOC数据集划分

‍uu分享的另一个好链接：DataAnalyze（对voc、coco等格式的数据集进行数据分析）

参考链接

• python将文件夹内的数据集按比例随机划分为训练集测试集,并移至其他目录
• yolov5训练数据集划分-爱代码爱编程

VOC数据格式

VOC数据集格式-参考

VOCdevkit(数据集名称)
├──	VOC2007
	    ├── Annotations 进行 detection 任务时的标签文件，xml 形式，文件名与图片名一一对应
	    ├── ImageSets 包含三个子文件夹 Layout、Main、Segmentation，其中 Main 存放的是分类和检测的数据集分割文件
	        ├── Main
	        │   ├── train.txt 写着用于训练的图片名称， 共 2501 个
	        │   ├── val.txt 写着用于验证的图片名称，共 2510 个
	        │   ├── trainval.txt train与val的合集。共 5011 个
	        │   ├── test.txt 写着用于测试的图片名称，共 4952 个
	    ├── JPEGImages 存放 .jpg 格式的图片文件
	    ├── SegmentationClass 存放按照 class 分割的图片      (object detection 不会使用到)
	    └── SegmentationObject 存放按照 object 分割的图片    (object detection 不会使用到)

代码

使用本代码方法：

1. 先确定一个数据集文件夹，先就命名为VOC2007吧（如上）
2. 在VOC2007文件夹下面手动构建 3 个子文件夹：
   Annotations：存放所有xml标签
   JPEGImages：存放所有图像
   ImageSets：运行本代码后，自动存储划分结果（将生成4个txt）
3. 主要有两个函数：split_imagesets_main 和 split_abspath
   • split_imagesets_main：划分图像名称
   • split_abspath：划分图像的绝对路径（参考下面main中的代码，第2点我注释掉了，有需要的，打开该部分代码即可）

注意：

• 需要自行指定正确文件地址：images_dir、xml_dir、save_dir，再运行本代码
• 我设置了随机数种子，能保证每次随机划分的结果相同
• 调用 img_and_existxml 函数，对图像目录中的文件进行判定：是否是图像文件、是否有xml标签（image_ext中指定的图像类型可自行更改）

import os
import random

image_ext = ['.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff']


def img_and_existxml(x, xml_list):  # 是图像文件且存在标签
    xname, ext = os.path.splitext(x)
    if ext not in image_ext:
        return False
    else:
        return xname + '.xml' in xml_list


def split_imagesets_main(images_dir, xml_dir, save_dir, rate1=0.9, rate2=0.9):
    """
    注意：以下3个文件夹最好都在统一根目录下，例如VOC2007
    :param images_dir: 存放所有原始图像
    :param xml_dir: 存放所有标签文件
    :param save_dir: 划分后的4个txt文件的存放目录，通常写为 .../VOC2007/ImageSets（此目录不存在不要紧，会自动生成）
    :param rate1: 训练集在总数据集中的比例，训练验证集：测试集=9：1（默认）
    :param rate2: 训练集在验证测试集中的比例，训练集：验证集=9：1（默认）
    :return: 无返回值，但最终会在.../VOC2007/ImageSets生成4个txt 只包含了文件名称，没有文件路径和后缀
    """
    assert os.path.exists(images_dir), f'\"{images_dir}\" not exists. split_imagesets_main over!'
    assert os.path.exists(xml_dir), f'\"{xml_dir}\" not exists. split_imagesets_main over!'

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    save_main_dir = os.path.join(save_dir, 'Main')
    os.makedirs(save_main_dir)

    xml_list = os.listdir(xml_dir)
    filelist = [x for x in os.listdir(images_dir) if img_and_existxml(x, xml_list)]

    random.seed(0)  # 设置随机数种子，保证每次得到的随机采样值相同
    trainval_list = random.sample(filelist, int(rate1 * len(filelist)))  # 随机采样
    train_list = random.sample(trainval_list, int(rate2 * len(trainval_list)))  # 随机采样

    trainval_txt = open(os.path.join(save_main_dir, 'trainval.txt'), 'w')
    train_txt = open(os.path.join(save_main_dir, 'train.txt'), 'w')
    val_txt = open(os.path.join(save_main_dir, 'val.txt'), 'w')
    test_txt = open(os.path.join(save_main_dir, 'test.txt'), 'w')

    num = [0, 0, 0, 0]
    res_txt = open(os.path.join(save_main_dir, 'result.txt'), 'w')

    for file in filelist:
        filename, _ = os.path.splitext(file)
        if file in trainval_list:
            trainval_txt.write(filename + '\n')  # 写入trainval.txt
            num[0] += 1
            if file in train_list:
                train_txt.write(filename + '\n')  # 写入train.txt
                num[1] += 1
            else:
                val_txt.write(filename + '\n')  # 写入val.txt
                num[2] += 1
        else:  # 测试集
            test_txt.write(filename + '\n')  # 写入test.txt
            num[3] += 1

    trainval_txt.close()
    train_txt.close()
    val_txt.close()
    test_txt.close()

    # 记录结果
    res_txt.write('trainval: {}'.format(num[0]) + '\n')
    res_txt.write('train: {}'.format(num[1]) + '\n')
    res_txt.write('val: {}'.format(num[2]) + '\n')
    res_txt.write('test: {}'.format(num[3]))
    res_txt.close()

    print('yes, saved in {}'.format(save_main_dir))
    print('trainval: {}'.format(num[0]))
    print('train: {}'.format(num[1]))
    print('val: {}'.format(num[2]))
    print('test: {}'.format(num[3]))


def split_abspath(images_dir, imagesets_main_dir, save_dir):
    """
    根据 调用split_imagesets_main函数 运行后生成的ImageSets/Main/train.txt val test trainval 里面划分的图像名称
    生成用于训练or验证or测试的图像的绝对路径
    :param images_dir: 原始图像的存放路径
    :param imagesets_main_dir: 存放已划分的txt文件目录，例如.../VOC2007/ImageSets
    :param save_dir: 任意指定的保存路径
    :return: 无返回值，但会在指定save_dir中生成txt记录相应图像的绝对路径
    """
    assert os.path.exists(images_dir), f'\"{images_dir}\" not exists. voc_labels over!'
    assert os.path.exists(imagesets_main_dir), f'\"{imagesets_main_dir}\" not exists. voc_labels over!'

    sets = ['train', 'val', 'test', 'trainval']
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    for s in sets:
        txt = os.path.join(imagesets_main_dir, s + '.txt')
        with open(txt, 'r') as f:
            file_list = f.read().strip().split()

        save_txt = os.path.join(save_dir, s + '.txt')
        with open(save_txt, 'w') as f:
            for file in file_list:
                f.write(os.path.join(images_dir, file + '.jpg\n'))  # 注意自己的图片格式是不是jpg


if __name__ == '__main__':

    images_dir = r'D:\A_dataset\threepeople\images'  # 所有图像路径
    xml_dir = r'D:\A_dataset\threepeople\labels_xml'  # 所有xml标签路径
    save_dir1 = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\Augment-lwd\caogao_2\ImageSets'

    # 1. 划分图像名称到txt中，rate1，rate2都使用默认值，使得训练集:验证集:测试集接近于8:1:1  (9:1):1
    split_imagesets_main(images_dir, xml_dir, save_dir1)

    # 2. 根据ImageSets/Main下的txt，将图像的绝对路径保存到到指定文件夹下的txt中
    # save_dir2 = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\Augment-lwd\caogao_2\save_2'
    # split_abspath(images_dir, save_dir1 + '/Main', save_dir2)