【yolo数据集】xml格式转换txt，划分数据集，亲测有效

项目场景：

在下载公共数据集时，有时得到的是xml格式的标签，我们需要转换成yolo所需要的txt格式，再按照需要的比例，划分数据集。在网上找了几个博主的方法，都出现了一些问题，最后做出来的txt文件存在各种问题，因此我综合多位博主的代码，自己改进了一份十分有效的转换方法。

---

解决方案：分为两段代码，请自行将数据集弄成voc数据集格式

 

 

 

 

下面一段代码，将所有的xml转换成了txt文件，保存在了自己新建的label目录下。只需要修改我提示修改的地方！！！

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import glob

classes = ["defect"]   #改为你的数据集类别   修改1


def convert(size, box):
    dw = 1.0 / (size[0]+1)  #这里加1是为了防止除数等于0
    dh = 1.0 / (size[1]+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)


def convert_annotation(image_name):
    in_file = open('defect/Annotations/' + image_name[:-3] + 'xml')  # xml文件路径   修改2
    out_file = open('./label/' + image_name[:-3] + 'txt', 'w')  # 转换后的txt文件存放路径  修改3   注意这里需要你自己新建一个目录label！！！
    f = open('defect/Annotations/' + image_name[:-3] + 'xml') # xml文件路径  修改4同修改1
    xml_text = f.read()
    root = ET.fromstring(xml_text)
    f.close()
    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
        if cls not in classes:
            print(cls)
            continue
        cls_id = classes.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')


wd = getcwd()

if __name__ == '__main__':

    for image_path in glob.glob("defect/JPEGImages/*.jpg"):  # 每一张图片都对应一个xml文件这里写xml对应的图片的路径   修改5
        image_name = image_path.split('/')[-1]
        print(image_name)
        convert_annotation(image_name)

下面是第二段代码，将图像和txt文件共同进行了切分

# 将一个文件夹下数据按比例分在三个文件夹下
import os
import random
import shutil
from shutil import copy2

datadir_normal = "label"  # 更换为自己刚刚转换出来的txt路径     修改1
datadir_normal1 = 'defect/JPEGImages' # 所有的图像路径          修改2

all_data = os.listdir(datadir_normal)  # （图片文件夹）
num_all_data = len(all_data)
print("总数量: " + str(num_all_data))
index_list = list(range(num_all_data))
print(index_list)

# 生成随机种子，使得每一次划分的数据一样，如果没有要求，可以删除本段代码
r = random.random
random.seed(2)
random.shuffle(index_list,random=r)
print(index_list)
# 生成随机种子，使得每一次划分的数据一样，如果没有要求，可以删除本段代码

num = 0
#下面设置你输出后的存放路径       注意新建两个目录，分别是labels和images！！！    修改3
#下面的目录不用修改，会自动生成！！！！！！！！！！
trainDir = "labels/train"  # （将训练集txt放在这个文件夹下）    不用修改
if not os.path.exists(trainDir):
    os.mkdir(trainDir)
validDir = 'labels/test'  # （将验证集txt放在这个文件夹下）     不用修改
if not os.path.exists(validDir):
    os.mkdir(validDir)
trainDir1 = 'images/train'  # （将训练集图片放在这个文件夹下）   不用修改
if not os.path.exists(trainDir1):
    os.mkdir(trainDir1)
validDir1 = 'images/test'  # （将验证集图片放在这个文件夹下）   不用修改
if not os.path.exists(validDir1):
    os.mkdir(validDir1)

for i in index_list:
    fileName = os.path.join(datadir_normal, all_data[i])
    if num < num_all_data * 0.8:  # Train:0.8
        print(str(fileName))
        copy2(fileName, trainDir)
    else:  # test:0.2
        copy2(fileName, validDir)
    num += 1
num = 0
for i in index_list:
    fileName = os.path.join(datadir_normal1, all_data[i].replace("txt","jpg"))
    if num < num_all_data * 0.8:
        copy2(fileName,trainDir1)
    else:  # test:0.2
        copy2(fileName, validDir1)
    num += 1

好啦，现在已经将数据集转换且划分完毕了，分别在images和labels里面，再把他们放到同一个文件夹中就可以在yolo里面使用了