python 修改xml文件中的目标类别名

在做目标检测的时候会遇到将两个或者多个类别看作一个类别进行检测,或者在标注的时候对类别名标注错误,需要对其进行修改。
下面是本人写的一个修改类别名的代码:

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET

import os


def convert_annotation(file, class_list, str_person):
    in_file = open(file, encoding='utf-8')
    tree = ET.parse(in_file)
    root = tree.getroot()
    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls in class_list:
            obj.find('name').text = str_person
            tree.write(in_file.buffer.name)
            print(in_file.buffer.name, "修改完毕")
    print("down")


if __name__ == '__main__':
    path = r''#xml文件路径
    class_list = ['people', 'cyclist', 'unpaired']#需要修改的类别名
    str_person = 'person'#修改后的类别名

    for file in os.listdir(path):
        image_id = file.split('.')[0]
        file = os.path.join(path, image_id + '.xml')
        convert_annotation(file, class_list, str_person)

结束

你可能感兴趣的:(python,xml,目标检测)