【python】修改目标检测的xml标签(VOC)类别名

需求：

在集成多个数据集一同训练时，可能会存在不同数据集针对同一种目标有不同的类名，可以通过python脚本修改数据内的类名映射，实现统一数据集标签名的目的。

代码：

# -*- coding: utf-8 -*-
# @Time : 2023/9/11 10:58
# @Author : CLW
# @FileName: change_xml_label.py
# @Software: PyCharm

import os
import xml.etree.ElementTree as ET

'''
算法功能：
针对目录下的所有voc格式的xml标签文件，将类别的名称修改为指定名称
要求：label_dict的key值表示修改前的标签名，该key对应的value值表示修改后的标签名
e.x. 需要将xml文件中'workwear'的名称替换为'reflectivevest'，那么可以这样设置：
    label_dict['workwear'] = 'reflectivevest'
'''

'''
####################    输入参数设置(开始)    #################### 
'''
org_label_dir = r'D:\dataset\pretrain_GW'
new_label_dir = r'D:\dataset\pretrain_GW'
label_dict = {}
label_dict['head'] = '头'
label_dict['hand'] = '手'


'''
####################    输入参数设置(结束)    #################### 
'''

def change_xml_label(xml_path, output_path, label_dict):
    if not os.path.getsize(xml_path):
        print("内容为空的xml:",xml_path)
    tree = ET.parse(xml_path)
    root = tree.getroot()
    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls in label_dict.keys():
            obj.find('name').text = label_dict[cls]

    tree.write(output_path, encoding="utf-8", xml_declaration=True)  # 输出新的xml文件

if __name__ == "__main__":
    # 遍历xml
    for root, dir, files in os.walk(org_label_dir):
        for file in files:
            if file[-3:] == 'xml':
                print(os.path.join(root, file))
                change_xml_label(os.path.join(root, file), os.path.join(new_label_dir, file),label_dict)