xml数据转yolov5的txt代码

通过labelimg标注的xml文件转换yolov5可训练的txt格式代码。包含读取xml文件代码，保存txt格式代码。

from lxml.etree import Element, SubElement, tostring, ElementTree
from xml.dom.minidom import parseString
import xml.etree.ElementTree as ET

import os

def get_root_lst(root, suffix='jpg', suffix_n=3):
    root_lst, name_lst = [], []

    for dir, file, names in os.walk(root):
        root_lst = root_lst + [os.path.join(dir, name) for name in names if name[-suffix_n:] == suffix]
        name_lst = name_lst + [name for name in names if name[-suffix_n:] == suffix]

    return root_lst, name_lst
def read_xml(xml_root):
    '''
    :param xml_root: .xml文件
    :return: dict('cat':['cat1',...],'bboxes':[[x1,y1,x2,y2],...],'whd':[w ,h,d])
    '''
    dict_info = {'cat': [], 'bboxes': [], 'box_wh': [], 'whd': []}
    if os.path.splitext(xml_root)[-1] == '.xml':
        tree = ET.parse(xml_root)  # ET是一个xml文件解析库，ET.parse（）打开xml文件。parse--"解析"
        root = tree.getroot()  # 获取根节点
        whd = root.find('size')
        whd = [whd.find('width').text, whd.find('height').text, whd.find('depth').text]

        for obj in root.findall('object'):  # 找到根节点下所有“object”节点
            cat = str(obj.find('name').text)  # 找到object节点下name子节点的值（字符串）
            bbox = obj.find('bndbox')
            x1, y1, x2, y2 = [int(bbox.find('xmin').text),
                              int(bbox.find('ymin').text),
                              int(bbox.find('xmax').text),
                              int(bbox.find('ymax').text)]
            b_w = x2 - x1 + 1
            b_h = y2 - y1 + 1

            dict_info['cat'].append(cat)
            dict_info['bboxes'].append([x1, y1, x2, y2])
            dict_info['box_wh'].append([b_w, b_h])
            dict_info['whd'].append(whd)
    else:
        print('[inexistence]:{} suffix is not xml '.format(xml_root))
    return dict_info
def build_dir(root):
    import os
    if not os.path.exists(root):
        os.makedirs(root)
    return root
# 保存txt格式文件
def write_txt(text_lst, out_txt=None):
    '''
    每行内容为列表，将其写入text中
    '''
    out_dir = out_txt if out_txt is not None else 'classes.txt'
    file_write_obj = open(out_dir, 'w', encoding='utf-8')  # 以写的方式打开文件，如果文件不存在，就会自动创建
    for text in text_lst:
        file_write_obj.writelines(str(text))
        file_write_obj.write('\n')
    file_write_obj.close()

    def get_str_name(root_str):
        if '\\' in root_str:
            root_str = root_str.split('\\')[-1]
        if '/' in root_str:
            root_str = root_str.split('/')[-1]
        return root_str
def xml2yolotxt(xml_root, img_root=None, out_dir=None, cat_name_lst=None):
    '''
    :param xml_root: xml的路径
    :param img_root:图像路径，可提供也可不提供，提供主要获得图像的高宽
    :param out_dir:保存txt路径
    :param cat_name_lst:提供训练列表，xml中出现类别与列表对应，如['pedes', 'elec', 'car', 'truck', 'bus', 'tricycle']
    pedes表示0，elec表示1，car表示2等
    :return:
    '''
    xml_root_lst, xml_names_lst = get_root_lst(xml_root, suffix='xml', suffix_n=3)
    out_dir = build_dir(out_dir) if out_dir is not None else build_dir(os.path.join(xml_root, 'out_dir_labels'))
    if img_root is not None:
        img_root_lst, img_names_lst = get_root_lst(img_root, suffix='jpg', suffix_n=3)

    for i, xml_root in tqdm(enumerate(xml_root_lst)):
        xml_info = read_xml(xml_root)
        if cat_name_lst is None:
            cat_lst = xml_info['cat']  # 类别是数字，从0 1 2 等
        else:
            cat_lst = [list(cat_name_lst).index(c) for c in xml_info['cat']]  # 类别名称，根据name列表得到数字类别

        if img_root is not None:
            # 从中提取W与H
            j = list(img_names_lst).index(xml_names_lst[i][:-3] + 'jpg')
            img = cv2.imread(img_root_lst[int(j)])
            H, W = img.shape[:2]
        else:
            whd = xml_info['whd'][0]
            W, H = float(whd[0]), float(whd[1])

        boxes_lst = xml_info['bboxes']
        yolotxt_lst = []
        for i, b in enumerate(boxes_lst):
            bw, bh = b[2] - b[0], b[3] - b[1]
            x, y = b[0] + bw / 2, b[1] + bh / 2
            x, y, w, h = x / W, y / H, bw / W, bh / H
            # yolotxt = str(cat_lst[i]) + ' ' + str(x) + ' ' + str(y) + ' ' + str(w) + ' ' + str(h)
            yolotxt = str(cat_lst[i]) + ' ' + str(x) + ' ' + str(y) + ' ' + str(w) + ' ' + str(h)

            yolotxt_lst.append(yolotxt)
        if len(yolotxt_lst) > 0:
            xml_name = get_str_name(xml_root)

            write_txt(yolotxt_lst, out_txt=os.path.join(out_dir, xml_name[:-3] + 'txt'))