使用python对label生成的xml及labelme生成的json数据集统计类别及类别的数量

有时需要检查校准训练数据集的类别及类别个数，需要对训练集进行统计，脚本对label生成的xml及labelme生成的json数据集统计类别及类别的个数。

import json
from glob import glob
import xml.etree.ElementTree as ET
import json
import os
labelme_path = '' #数据集路径

def get_and_check(root, name, length):
    vars = root.findall(name)
    if len(vars) == 0:
        raise NotImplementedError('Can not fing %s in %s.' % (name, root.tag))
        if length > 0 and len(vars) != length:
             raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))
        if length == 1:
        vars = vars[0]
    return vars
def write_xml(tree, out_path):
  tree.write(out_path, encoding="utf-8", xml_declaration=True)
def get_labels(ann_dir, metric_type):
     """
    获得标签
    :param ann_dir: 标签路径
    :param metric_type: 评价类型
    :return:
    """
    ann_path_list = os.listdir(ann_dir)
    label_num_dict = {}
    if metric_type == "bbox":
        for xml_i in ann_path_list:
            if 'xml' in xml_i:
                file_path = os.path.join(ann_dir, xml_i)
                tree = ET.parse(file_path)
                root = tree.getroot()
                for obj in root.findall('object'):
                    category = get_and_check(obj, 'name', 1).text
                     if category not in label_num_dict.keys():
                        num = 1
                        label_num_dict[category] = num
                     else:
                        for key, value in label_num_dict.items():
                            if key == category:
                                label_num_dict[key] = value + 1
    elif metric_type == "segm":
         for json_i in ann_path_list:
            if 'json' in json_i:
                with open(os.path.join(ann_dir, json_i), "r", encoding='utf-8') as f:
                     obj = json.load(f)
                shapes = obj['shapes']
                for shape in shapes:
                    label = shape['label']
                    if label not in label_num_dict.keys():
                        num = 1
                        label_num_dict[category] = num
                     else:
                        for key, value in label_num_dict.items():
                            if key == category:
                                label_num_dict[key] = value + 1
    return label_num_dict
label = get_labels(labelme_path, "bbox") #调用