目标检测小脚本：根据xml文件统计类别数

Posted 2022-08-18 zstar-_

问题场景

搜到一个目标检测数据集，但是别人没有详细说明到底有多少类别以及各类别名称，这时候就需要查询xml的标注文件来获取这两个信息。

脚本代码

import xml.etree.ElementTree as ET
import numpy as np
import os

if __name__ == '__main__':
    xmlfilepath = 'D:/Dataset/mydata/Annotations'  # xml文件存放地址
    classes = []
    for i in os.listdir(xmlfilepath):
        in_file = open(xmlfilepath + '/' + i)
        tree = ET.parse(in_file)
        root = tree.getroot()
        for obj in root.iter('object'):
            cls = obj.find('name').text
            classes.append(cls)
    print("总类别数:".format(len(np.unique(classes))))
    print(np.unique(classes))