python实用脚本（二）：删除xml标签下的指定类别

CV-King2024-10-01 13:52

介绍

在目标检测中，有些时候会遇到标注好的类别不想要了的情况，这时我们可以运行下面的代码来批量删除不需要的类别节省时间。

代码实现：

python 复制代码

import argparse

import xml.etree.ElementTree as ET
import os

classes = ['thin_smoke']


def GetImgNameByEveryDir(file_dir, videoProperty):
    FileNameWithPath, FileName, FileDir = [], [], []
    for root, dirs, files in os.walk(file_dir):
        for file in files:
            if os.path.splitext(file)[1] in videoProperty:
                FileNameWithPath.append(os.path.join(root, file))  # 保存图片路径
                FileName.append(file)  # 保存图片名称
                FileDir.append(root[len(file_dir):])  # 保存图片所在文件夹
    return FileName, FileNameWithPath, FileDir


def GetBoxInfo(xmlfile):
    try:
        tree = ET.parse(xmlfile)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
    except:
        return False, 0
    else:
        tree = ET.parse(xmlfile)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)

    for obj in root.findall('object'):
        print(xmlfile, obj.find('name').text, obj)
        if obj.find('name').text == 'fulll':
            obj.find('name').text = "full"

    print('-' * 66)
    tree.write(xmlfile)

    return 1


def Process(ProcessDir):
    xmlDirs = ProcessDir  # + 'Labels/'
    FileName1, FileNameWithPath1, FileDir1 = GetImgNameByEveryDir(xmlDirs, '.xml')

    for k in range(len(FileName1)):
        annfile = xmlDirs + FileName1[k][:-4] + '.xml'
        result = GetBoxInfo(annfile)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--ProcessDir', type=str, default='/home/build/smoke_ori/VOCdevkit/VOC2007/Annotations/')
    args = parser.parse_args()
    ProcessDir = args.ProcessDir
    Process(ProcessDir)