数据集标签文件转换方法--将 XML 文件类型转化为 TXT 文件类型

数据集标签文件转换方法：
数据标注软件labelimg，标注VOC格式数据后缀为xml，YOLO格式为txt文件。
image_split_to_txt.py文件, 用于将 XML 文件类型转化为 TXT 文件类型并将对应的图片和 XML 文件移动到训练、验证、测试集的目录中.
注意：
每次重新生成时，需要将原文件删掉，否则不更新。
代码里classes，改成自己的类型。
每次重新生成时，原始的数据集要重新复制，因为每次转换后，原文件夹中用到的图片移出了。
python 复制代码
"""
将xml标签文件转换为txt文件。注意：代码中类别标签一定要修改
"""
# 1.导入相应的库
import os
import random
import argparse
import shutil
import xml.etree.ElementTree as ET
from tqdm import tqdm
 
# 参数解析
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='../data/dataset', type=str, help='input xml label path')
parser.add_argument('--images_path', default='../data/images', type=str, help='original images path')
parser.add_argument('--output_path', default='../data', type=str, help='output base path')
opt = parser.parse_args()
 
# 定义划分比例
trainval_percent = 1.0
train_percent = 0.8
val_percent = 0.2
 
# 设置路径
xml_path = opt.xml_path
images_path = opt.images_path
output_path = opt.output_path
 
# 创建输出目录结构
os.makedirs(os.path.join(output_path, 'images/train'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'images/val'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'images/test'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'labels/train'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'labels/val'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'labels/test'), exist_ok=True)
 
# 获取所有XML文件
xml_files = [f for f in os.listdir(xml_path) if f.endswith('.xml')]
random.shuffle(xml_files)
num = len(xml_files)
 
# 划分数据集
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
ta = tv - tr  # 确保train+val=trainval
 
trainval = xml_files[:tv]
train = trainval[:tr]
val = trainval[tr:tr + ta]
test = xml_files[tv:]
 
# 类别定义
classes = ["panda","no_panda"]
 
# XML转换函数
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    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, output_dir):
    try:
        image_id = os.path.splitext(xml_file)[0]
        in_file = os.path.join(xml_path, xml_file)
        out_file = os.path.join(output_dir, f'{image_id}.txt')
 
        tree = ET.parse(in_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
 
        with open(out_file, 'w', encoding='utf-8') as f_out:
            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
                cls_id = classes.index(cls)
                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((w, h), b)
                f_out.write(f"{cls_id} {' '.join(str(a) for a in bb)}\n")
 
        return image_id
    except Exception as e:
        print(f"Error processing {xml_file}: {str(e)}")
        return None
 
 
# 处理训练集
print("Processing training set...")
train_list = []
for xml_file in tqdm(train):
    image_id = convert_annotation(xml_file, os.path.join(output_path, 'labels/train'))
    if image_id:
        src_img = os.path.join(images_path, f'{image_id}.jpg')
        dst_img = os.path.join(output_path, 'images/train', f'{image_id}.jpg')
        if os.path.exists(src_img):
            # 移动图片文件到训练集目录
            shutil.move(src_img, dst_img)
            train_list.append(dst_img)
 
        # 移动XML文件到标签目录
        src_xml = os.path.join(xml_path, xml_file)
        dst_xml = os.path.join(output_path, 'labels/train', f'{image_id}.xml')
        if os.path.exists(src_xml):
            shutil.move(src_xml, dst_xml)
 
# 处理验证集
print("Processing validation set...")
val_list = []
for xml_file in tqdm(val):
    image_id = convert_annotation(xml_file, os.path.join(output_path, 'labels/val'))
    if image_id:
        src_img = os.path.join(images_path, f'{image_id}.jpg')
        dst_img = os.path.join(output_path, 'images/val', f'{image_id}.jpg')
        if os.path.exists(src_img):
            # 移动图片文件到验证集目录
            shutil.move(src_img, dst_img)
            val_list.append(dst_img)
 
        # 移动XML文件到标签目录
        src_xml = os.path.join(xml_path, xml_file)
        dst_xml = os.path.join(output_path, 'labels/val', f'{image_id}.xml')
        if os.path.exists(src_xml):
            shutil.move(src_xml, dst_xml)
 
# 处理测试集
print("Processing test set...")
test_list = []
for xml_file in tqdm(test):
    image_id = convert_annotation(xml_file, os.path.join(output_path, 'labels/test'))
    if image_id:
        src_img = os.path.join(images_path, f'{image_id}.jpg')
        dst_img = os.path.join(output_path, 'images/test', f'{image_id}.jpg')
        if os.path.exists(src_img):
            # 移动图片文件到测试集目录
            shutil.move(src_img, dst_img)
            test_list.append(dst_img)
 
        # 移动XML文件到标签目录
        src_xml = os.path.join(xml_path, xml_file)
        dst_xml = os.path.join(output_path, 'labels/test', f'{image_id}.xml')
        if os.path.exists(src_xml):
            shutil.move(src_xml, dst_xml)
 
 
# 保存路径文件
def save_path_file(file_path, path_list):
    with open(file_path, 'w') as f:
        for path in path_list:
            f.write(f"{path}\n")
 
 
save_path_file(os.path.join(output_path, 'train.txt'), train_list)
save_path_file(os.path.join(output_path, 'val.txt'), val_list)
save_path_file(os.path.join(output_path, 'test.txt'), test_list)
 
print("Dataset preparation completed!")
print(f"Train: {len(train_list)} images, Val: {len(val_list)} images, Test: {len(test_list)} images")