修改xml文件标签名称、筛选xml文件中的指定标签

在模型训练和数据标注的时候经常出现一些需要修改标签信息的情况，或者说想要单独整理一类的标签出来审核或者训练，通常我们用labelimg来一张张图片审核修改是非常吃力的，所以我写了一个脚本方便大家对标签信息做一些自动化操作。请注意使用的时候需要在要保存的文件夹路径下创建名为img和xml的空文件夹
这篇文章主要解决两个需求：
1、在整个数据集中将一个标签集合的图片和标注筛除出来，比如在大数据集中有ABC三个类，如果此时我想要将B类的图片和标注单独分出来并且选出来的标注中会删除A和C，可以使用下面这个脚本：

# 从原始img和xml中筛选拥有目标label的img和xml到新路径下，并且新xml中只保留我要的标签
import os
import shutil
from pathlib import Path
import xml.etree.ElementTree as ET

output_dir = r'new_path'  # 筛选修改后的新xml保存路径
input_dir = r'raw_path'  # 原始的xml文件夹

# 可以自定义自己想要筛选的标签, 
# large_labels = ['A', 'B', 'C']
target_labels = ['B']. # 只需要打标签集合的子集就行了，比如这里就只会将带有B标签的图片选出来并且对应的标注中不会有A、C的标签信息

# 初始化路径
input_dir_path = Path(input_dir)
output_dir_path = Path(output_dir)
new_xml_dir = output_dir_path / 'xml'
raw_xml_dir = input_dir_path / 'xml'
new_img_dir = output_dir_path / 'img'
raw_img_dir = input_dir_path / 'img'

n = 0  # 统计筛选的数量
img_names = []  # 记录筛选出来的图片名称
for xml in os.listdir(raw_xml_dir):
    raw_xml_path = str(raw_xml_dir / xml)
    new_xml_path = str(new_xml_dir / xml)  # 后面保存修改后的xml文件时用到
    tree = ET.parse(raw_xml_path)
    root = tree.getroot()
    print(f'处理{raw_xml_path}中')
    for obj in root.findall('object'):
        name = obj.find('name').text
        if name not in target_labels:
            root.remove(obj)
    # 检查标签是否全部删掉，如果标签集不为空则保存
    if root.findall('object'):
        n += 1
        img_names.append(xml[:-4])
        tree.write(new_xml_path)

# 复制对应图片到新路径下
for img in os.listdir(raw_img_dir):
    raw_img_path = str(raw_img_dir / img)
    new_img_path = str(new_img_dir / img)
    if img[:-4] in img_names:
        shutil.copy(raw_img_path, new_img_path)
print('图片拷贝完成。')
print(f'总计挑选出{n}组img和xml。')

2、如果我在训练的时候发现有标注中有标签名写错了导致训练报错或者说我想要在原始标签文件中将一类标签名进行修改可以使用下面的脚本：

# 从原始img和xml中筛选拥有目标label的img和xml到新路径下，并且新xml中修改标签名称
import os
import shutil
from pathlib import Path
import xml.etree.ElementTree as ET

output_dir = r'new_path'  # 筛选修改后的新xml保存路径
input_dir = r'raw_path'  # 原始的xml文件夹

# 新旧标签名称映射字典
label_pair = {'A': 'B'}. # 将A标签改成B

# 初始化路径
input_dir_path = Path(input_dir)
output_dir_path = Path(output_dir)
new_xml_dir = output_dir_path / 'xml'
raw_xml_dir = input_dir_path / 'xml'
new_img_dir = output_dir_path / 'img'
raw_img_dir = input_dir_path / 'img'

n = 0  # 统计筛选的数量
img_names = []  # 记录筛选出来的图片名称
for xml in os.listdir(raw_xml_dir):
    raw_xml_path = str(raw_xml_dir / xml)
    new_xml_path = str(new_xml_dir / xml)  # 后面保存修改后的xml文件时用到
    tree = ET.parse(raw_xml_path)
    root = tree.getroot()
    print(f'处理{raw_xml_path}中')
    for obj in root.findall('object'):
        name = obj.find('name').text
        if name in label_pair:
            obj.find('name').text = label_pair[name]
    # 如果标签集不为空则保存
    if root.findall('object'):
        n += 1
        img_names.append(xml[:-4])
        tree.write(new_xml_path)

# 复制对应图片到新路径下
for img in os.listdir(raw_img_dir):
    raw_img_path = str(raw_img_dir / img)
    new_img_path = str(new_img_dir / img)
    if img[:-4] in img_names:
        shutil.copy(raw_img_path, new_img_path)
print('图片拷贝完成。')
print(f'总计处理{n}组img和xml。')

注意：我这里为了防止误操作都是将数据修改之后保存到一个新的路径而不是替换原来的文件，如果需要可以自己修改一下代码。