labelimg标注的xml标签转换为yolo格式标签

本文不生产技术，只做技术的搬运工！！！

前言

        在yolo训练时，我们需要对图像进行标注，而使用labelimg标注时如果直接选择输出yolo格式的数据集，则原始数据的很多信息无法被保存，因此一版使用xml格式的标签，这时再去训练时就需要对标签进行转换。

代码

import os
import xml.etree.ElementTree as ET
import cv2

def getFileList(dir, Filelist, ext=None):
    """
    获取文件夹及其子文件夹中文件列表
    输入 dir：文件夹根目录
    输入 ext: 扩展名
    返回： 文件路径列表
    """
    newDir = dir
    if os.path.isfile(dir):
        if ext is None:
            Filelist.append(dir)
        else:
            if ext in dir:
                Filelist.append(dir)

    elif os.path.isdir(dir):
        for s in os.listdir(dir):
            newDir = os.path.join(dir, s)
            getFileList(newDir, Filelist, ext)

    return Filelist

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)

def convert_annotation(xml_file,xml_name, image_file,dst, class_names):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 读取图像以获取准确的尺寸
    image = cv2.imread(image_file)
    if image is None:
        raise FileNotFoundError(f"Image file not found: {image_file}")
    image_size = [image.shape[1], image.shape[0]]  # [width, height]

    output_path = os.path.join(dst, xml_name.replace(".xml", '.txt'))

    with open(output_path, 'w') as out_file:
        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            class_name = obj.find('name').text
            if class_name not in class_names or int(difficult) == 1:
                continue
            class_id = class_names.index(class_name)
            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(image_size, b)
            out_file.write(f"{class_id} {bb[0]} {bb[1]} {bb[2]} {bb[3]}\n")

if __name__ == '__main__':
    xml_dir = "/data2/dataset/abcd/20250305/xml" #xml标签路径
    image_dir = "/data2/dataset/abcd/20250305/images" #图像路径
    output_dir = "/data2/dataset/abcd/20250305/labels" #输出的txt路径
    class_names = ['aaa', 'bbb', 'ccc', 'ddd']  # 类别名称列表
    xml_file_list = []
    xml_file_list = getFileList(xml_dir, xml_file_list, '.xml')
    lenth = len(xml_file_list)
    i = 1
    for xml_file in xml_file_list:
        xml_name = os.path.basename(xml_file)
        image_name = xml_name.replace(".xml", ".jpg")
        image_file = os.path.join(image_dir, image_name)
        if not os.path.exists(image_file):
            print(f"image file not found: {image_file}")
            continue
        convert_annotation(xml_file, xml_name, image_file, output_dir, class_names)
        print('{}/{}'.format(i, lenth))
        i += 1