使用开源数据集 遇到的问题，代码：.

import os
import xml.etree.ElementTree as ET

# 定义类别及其对应的 class_id
class_mapping = {
    'crazing': 0,
    'inclusion': 1,
    'patches': 2,
    'pitted_surface': 3,
    'rolled-in_scale': 4,
    'scratches': 5,
    # 添加其他类别和对应的ID
}

def convert_xml_to_yolo(xml_file, output_dir):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 获取图像宽度和高度
    width = int(root.find('size/width').text)
    height = int(root.find('size/height').text)

    # 创建 YOLO 格式的标注文件
    yolo_lines = []
    for obj in root.findall('object'):
        class_name = obj.find('name').text
        if class_name not in class_mapping:
            continue  # 忽略未定义的类别

        class_id = class_mapping[class_name]

        # 获取边界框坐标
        xmin = int(obj.find('bndbox/xmin').text)
        ymin = int(obj.find('bndbox/ymin').text)
        xmax = int(obj.find('bndbox/xmax').text)
        ymax = int(obj.find('bndbox/ymax').text)

        # 转换为 YOLO 格式
        x_center = (xmin + xmax) / 2 / width
        y_center = (ymin + ymax) / 2 / height
        bbox_width = (xmax - xmin) / width
        bbox_height = (ymax - ymin) / height

        # 保存一行标注
        yolo_lines.append(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}\n")

    # 输出 YOLO 格式文件
    output_file = os.path.join(output_dir, os.path.splitext(os.path.basename(xml_file))[0] + ".txt")
    with open(output_file, "w") as f:
        f.writelines(yolo_lines)

# 设置输入 XML 文件夹和输出 YOLO 文件夹
xml_dir = "/home/NEU-DET/ANNOTATIONS"
output_dir = "/home/NEU-DET/labels"
os.makedirs(output_dir, exist_ok=True)

# 批量转换
for xml_file in os.listdir(xml_dir):
    if xml_file.endswith(".xml"):
        convert_xml_to_yolo(os.path.join(xml_dir, xml_file), output_dir)