xml转txt

"""
作者: stark
日期: 2024年 国庆前夕
根据XML文件中的标签名自动找到相应的图像文件，并根据实际图像尺寸将XML格式转换为YOLO对应的txt格式。这个脚本假设XML文件中的图像文件名与实际图像文件名相对应，并且图像文件存储在指定的文件夹中。
直接运行就可以，运行后，输入一个主路径（存放有标签和图片）就可以。会自动在该路径下生成output_txt文件夹，里面就是转换好的txt文件。
"""
import xml.etree.ElementTree as ET
import os
from PIL import Image

def find_image(filename, base_dir):
    # 在整个base_dir路径下查找图像文件
    for root, dirs, files in os.walk(base_dir):
        if filename in files:
            return os.path.join(root, filename)
    return None

def convert_voc_to_yolo(xml_file, output_dir, class_mapping, img_dir):
    # 解析XML文件
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 获取图像的文件名
    filename = root.find('filename').text

    # 在指定的图像文件夹或整个目录中查找图像
    img_path = find_image(filename, img_dir)
    if not img_path:
        print(f"找不到图像文件：{filename}")
        return

    # 打开图像，获取宽度和高度
    try:
        with Image.open(img_path) as img:
            img_width, img_height = img.size
    except FileNotFoundError:
        print(f"找不到图像文件：{img_path}")
        return

    # 创建输出文件的TXT路径
    txt_file = os.path.join(output_dir, os.path.splitext(filename)[0] + '.txt')

    with open(txt_file, 'w') as f_out:
        # 遍历所有object节点
        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]
            bndbox = obj.find('bndbox')

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

            # 计算 YOLO 格式中的 (center_x, center_y, width, height)，并归一化
            center_x = (xmin + xmax) / 2.0 / img_width
            center_y = (ymin + ymax) / 2.0 / img_height
            width = (xmax - xmin) / img_width
            height = (ymax - ymin) / img_height

            # 写入YOLO格式的标签
            f_out.write(f"{class_id} {center_x:.6f} {center_y:.6f} {width:.6f} {height:.6f}\n")

    print(f"转换完成：{txt_file}")

if __name__ == "__main__":
    # 你的类别映射 (VOC类名: YOLO类ID)
    class_mapping = {
        'person': 0,
        'car': 1,
        'bicycle': 2,
        # 根据你的数据集添加更多类别
    }

    # 输入路径
    base_dir = input("请输入主目录路径：")

    output_dir = os.path.join(base_dir, "output_txt")
    img_dir = base_dir  # 假设图像文件和XML文件在同一个主目录或其子目录中

    os.makedirs(output_dir, exist_ok=True)

    # 递归查找所有XML文件
    for root, dirs, files in os.walk(base_dir):
        for file in files:
            if file.endswith(".xml"):
                xml_path = os.path.join(root, file)
                convert_voc_to_yolo(xml_path, output_dir, class_mapping, img_dir)