有时需要检查校准训练数据集的类别及类别个数,需要对训练集进行统计,脚本对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") #调用