统计Pascal VOC格式数据集中每一个类别Bounding box大小

import os
import xml.etree.ElementTree as ET
from collections import defaultdict

def get_annotations_info(file_path):
    tree = ET.parse(file_path)
    root = tree.getroot()

    info = []
    for obj in root.iter('object'):
        name = obj.find('name').text
        bbox = obj.find('bndbox')
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)

        width = xmax - xmin
        height = ymax - ymin

        info.append((name, width, height))
    return info

def calculate_avg_bbox_size(directory):
    stats = defaultdict(lambda: [0, 0, 0])  # Format: classname -> [total_width, total_height, count]

    for file_name in os.listdir(directory):
        if not file_name.endswith('.xml'):
            continue

        file_path = os.path.join(directory, file_name)
        annotations = get_annotations_info(file_path)

        for name, width, height in annotations:
            stats[name][0] += width
            stats[name][1] += height
            stats[name][2] += 1

    avg_sizes = {name: [total_width / count, total_height / count] for name, (total_width, total_height, count) in stats.items()}

    return avg_sizes

# Replace with your directory path
directory = "./Annotations/"
avg_sizes = calculate_avg_bbox_size(directory)

for name, (avg_width, avg_height) in avg_sizes.items():
    print(f'Class {name}: average width = {avg_width}, average height = {avg_height}')