VOC类型数据集划分为训练集、验证集、测试集

对于labelme采集的VOC类型数据集。我们一般有将其划分为训练集、验证集、 测试集方便我们模型使用的需求。

针对上述问题，我编写了一个小工具代码。

import os
import glob
import random
import xml.etree.ElementTree as ET
config = {
    # Annotations path(Annotations 的文件夹路径)
    "Annotation":"/data/data94796/Annotations",
    # JPEGImages path(JPEGImages 的文件夹路径)
    "JPEGImages":"/data/data94796/images",
}
# 划分数据集

# 数据划分比例
# (训练集+验证集)与测试集的比例，默认情况下 (训练集+验证集):测试集 = 9:1

# 按照比例划分数据集
train_per = 0.8
valid_per = 0.1
test_per = 0.1

data_xml_list = glob.glob(os.path.join(config['Annotation'], '*.xml'))
data_jpg_list = glob.glob(os.path.join(config['JPEGImages'], '*.*'))
data_xml_list.sort()
data_jpg_list.sort()
# 生成label标签:
label = set()
for xml_path in data_xml_list:
        label = label | set([i.find('name').text for i in ET.parse(xml_path).findall('object')])
data_list=[]
for i,j in zip(data_jpg_list,data_xml_list):
    data_list.append(i+" "+j)
data_xml_list=data_list
print(data_xml_list)
random.seed(666)
random.shuffle(data_xml_list)
data_length = len(data_xml_list)

train_point = int(data_length * train_per)
train_valid_point = int(data_length * (train_per + valid_per))

# 生成训练集，验证集, 测试集(8 : 1 : 1)
train_list = data_xml_list[:train_point]
valid_list = data_xml_list[train_point:train_valid_point]
test_list = data_xml_list[train_valid_point:]

 


# 写入文件中
ftrain = open('/home/aistudio/data/data94796/trainval.txt', 'w')
fvalid = open('/home/aistudio/data/data94796/valid.txt', 'w')
ftest = open('/home/aistudio/data/data94796/test.txt', 'w')
flabel = open('/home/aistudio/data/data94796/label_list.txt', 'w')

for i in train_list:
        ftrain.write(i + "\n")
for j in valid_list:
        fvalid.write(j + "\n")
for k in test_list:
        ftest.write(k + "\n")
for l in label:
        flabel.write(l + "\n")
ftrain.close()
fvalid.close()
ftest.close()
flabel.close()
print("总数据量:{}, 训练集:{}, 验证集:{}, 测试集:{}, 标签:{}".format(len(data_xml_list), len(train_list), len(valid_list), len(test_list), len(label)))
print("done!")

效果如下。左边会生成label标签文件，以及划分的训练集, 验证集， 测试集文件。
这里文件内的组织形式的 图片地址+标签地址。

若只需要标签地址可以使用这个代码。

import os
import glob
import random
import xml.etree.ElementTree as ET
config = {
    # Annotations path(Annotations 的文件夹路径)
    "Annotation":"/data/data94796/Annotations",
    # JPEGImages path(JPEGImages 的文件夹路径)
    "JPEGImages":"/data/data94796/images",
}
# 划分数据集

# 数据划分比例
# (训练集+验证集)与测试集的比例，默认情况下 (训练集+验证集):测试集 = 9:1

# 按照比例划分数据集
train_per = 0.8
valid_per = 0.1
test_per = 0.1

data_xml_list = glob.glob(os.path.join(config['Annotation'], '*.xml'))
# 生成label标签:
label = set()
for xml_path in data_xml_list:
        label = label | set([i.find('name').text for i in ET.parse(xml_path).findall('object')])
random.seed(666)
random.shuffle(data_xml_list)
data_length = len(data_xml_list)

train_point = int(data_length * train_per)
train_valid_point = int(data_length * (train_per + valid_per))

# 生成训练集，验证集, 测试集(8 : 1 : 1)
train_list = data_xml_list[:train_point]
valid_list = data_xml_list[train_point:train_valid_point]
test_list = data_xml_list[train_valid_point:]

 


# 写入文件中
ftrain = open('/home/aistudio/data/data94796/trainval.txt', 'w')
fvalid = open('/home/aistudio/data/data94796/valid.txt', 'w')
ftest = open('/home/aistudio/data/data94796/test.txt', 'w')
flabel = open('/home/aistudio/data/data94796/label_list.txt', 'w')

for i in train_list:
        ftrain.write(i + "\n")
for j in valid_list:
        fvalid.write(j + "\n")
for k in test_list:
        ftest.write(k + "\n")
for l in label:
        flabel.write(l + "\n")
ftrain.close()
fvalid.close()
ftest.close()
flabel.close()
print("总数据量:{}, 训练集:{}, 验证集:{}, 测试集:{}, 标签:{}".format(len(data_xml_list), len(train_list), len(valid_list), len(test_list), len(label)))
print("done!")