VOC格式标签转YOLO格式

VOC转Yolo格式代码:​​​​​​​

import os
import xml.etree.ElementTree as ET


# 将x1, y1, x2, y2转换成yolo所需要的x, y, w, h格式
def voc_to_yolo_wh(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[2]) / 2 * dw
    y = (box[1] + box[3]) / 2 * dh
    w = (box[2] - box[0]) * dw
    h = (box[3] - box[1]) * dh
    return x, y, w, h  # 返回的都是标准化后的值


# 获取所以xml里面的classes类别,返回(列表/数组)
def get_cls_list(xml_path):
    result_list = []
    for file in os.listdir(xml_path):
        tree = ET.parse(xml_path + file)
        root = tree.getroot()
        for obj in root.iter('object'):
            cls = obj.find('name').text

            if cls not in result_list:
                result_list.append(cls)
    return result_list


def voc_to_yolo(xml_path, save_path, classes):
    # 可以打印看看该路径是否正确

    if not os.path.exists(xml_path):
        print("XML路径不存在")
        return []

    if len(os.listdir(xml_path)) == 0:
        print("XML文件目录为空")
        return []

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

    if not classes:
        classes = get_cls_list(xml_path)

    # 遍历每一个xml文件
    for file in os.listdir(xml_path):
        label_file = xml_path + file
        save_file = save_path + file.replace('xml', 'txt')
        out_file = open(save_file, 'w')
        # print(label_file)

        # 开始解析xml文件
        tree = ET.parse(label_file)
        root = tree.getroot()
        size = root.find('size')  # 图片的shape值
        w = int(size.find('width').text)
        h = int(size.find('height').text)

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                continue
            # 将名称转换为id下标
            cls_id = classes.index(cls)
            # 获取整个bounding box框
            bnd_box = obj.find('bndbox')
            # xml给出的是x1, y1, x2, y2
            box = [float(bnd_box.find('xmin').text), float(bnd_box.find('ymin').text), float(bnd_box.find('xmax').text),
                   float(bnd_box.find('ymax').text)]

            # 将x1, y1, x2, y2转换成yolo所需要的x_center, y_center, w, h格式
            bbox = voc_to_yolo_wh((w, h), box)

            # print(save_file)
            # 写入目标文件中，格式为 id x y w h
            out_file.write(str(cls_id) + " " + " ".join(str(x) for x in bbox) + '\n')
    return classes


if __name__ == '__main__':
    # 这里要改成自己数据集路径
    xml_path_ = 'Annotations/xml/'
    # 这里要改成转换后保存的数据集路径
    save_path_ = 'Annotations/labels/'
    # 数据集标签,为空则自动获取标签,只需要在函数返回结果里获取标签就可以了
    Class_Name = []
    result_classes = voc_to_yolo(xml_path_, save_path_, Class_Name)
    print("数据集标签:", result_classes)