数据集格式转换vocToYolol

简述

针对VOC数据集转换为YOLO总共分为两个部分，首先对标签文件进行转换，进而对数据集划分训练与验证集。

标签文件转换

import xml.etree.ElementTree as ET
import os

voc_folder = r"VOC2007/Annotations"  # 储存voc格式的标注文件的文件夹
yolo_folder = r"VOC2007/Annotations_yolo"  # 转换后的yolo格式标注文件的储存文件夹

# 检查文件夹是否存在
assert os.path.exists(voc_folder), "voc_folder path not exist"
if os.path.exists(yolo_folder) is False:
    os.makedirs(yolo_folder)


class_id = ["hat", "person"]  # 储存数据集中目标种类名称的列表，接下来的转换函数中会将该列表中种类名称对应的列表索引号作为写入yolo标注文件中该类目标的种类序号


# voc标注的目标框坐标值转换到yolo标注的目标框坐标值的函数
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


# 对单个voc标注文件进行转换成其对应的yolo文件的函数
def convert_annotation(xml_file):
    file_name = xml_file.strip(".xml")  # 这一步将所有voc格式标注文件取出后缀名“.xml”，方便接下来作为yolo格式标注文件的名称
    in_file = open(os.path.join(voc_folder, xml_file), 'r', encoding='utf-8')  # 打开当前转换的voc标注文件
    out_file = open(os.path.join(yolo_folder, file_name + ".txt", ), 'w')  # 创建并打开要转换成的yolo格式标注文件
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        cls = obj.find('name').text
        cls_id = class_id.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text),
             float(xmlbox.find('xmax').text),
             float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


xml_fileList = os.listdir(voc_folder)  # 将所有voc格式的标注文件的名称取出存放到列表xml_fileList中
for xml_file in xml_fileList:  # 这里的for循环开始依次对所有voc格式标注文件生成其对应的yolo格式的标注文件
    convert_annotation(xml_file)

数据集划分

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
# 数据集的地址，根据自己的数据进行修改
parser.add_argument('--xml_path', default='VOC2007', type=str, help='input xml label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='VOC2007/ImageSets/Main', type=str, help='output train.txt path')
# 标签文件存放地址
parser.add_argument('--label_path',default='VOC2007/Annotations_yolo', type=str, help='label.txt path')
opt = parser.parse_args()

trainval_percent = 1  # 训练集和验证集所占比例。 这里没有划分测试集
train_percent = 0.9     # 训练集所占比例，可自己进行调整
'''
由于没有分配测试集，所以测试集为空。
若要分配，更改第 14、15 行代码，更改所在比例即可。
'''

datasetsPath = opt.xml_path
txtsavepath = opt.txt_path
labelPath = opt.label_path
total_txt = os.listdir(labelPath)

num = len(total_txt)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    # name = 'data/datasets/images/' + total_txt[i][:-4] + '.jpg' '\n'
    name = '%s/JPEGImages/%s.jpg'%(os.path.abspath(datasetsPath), total_txt[i][:-4]) + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()