统计coco数据集标签数量

一、统计coco数据集标签数量

在目标检测任务中，了解数据集中各个类别的数量是非常重要的。通过统计类别数量，可以了解数据集的分布情况，进而为模型训练和评估提供参考。

本文将介绍如何使用Python编写一个程序，来统计目标检测数据集中各个类别的数量。将使用Python的os和json库来读取和处理JSON格式的注释文件。

二、关键代码

import os
import json

def count_category_ids(folder_path, target_id=None):
    # 获取文件夹中所有的JSON文件
    json_files = [f for f in os.listdir(folder_path) if f.endswith('coco_filtered.json')]

    category_id_counts = {}  # 存储各个 category_id 的数量

    # 循环遍历每个JSON文件
    for file_name in json_files:
        file_path = os.path.join(folder_path, file_name)

        # 读取JSON文件
        with open(file_path, 'r') as file:
            data = json.load(file)

        # 统计各个 category_id 的数量
        for annotation in data['annotations']:
            category_id = annotation.get('category_id')

            if category_id is not None:
                if target_id is None or category_id == target_id:
                    if category_id not in category_id_counts:
                        category_id_counts[category_id] = 0
                    category_id_counts[category_id] += 1

    # 打印各个 category_id 的数量
    for category_id, count in category_id_counts.items():
        print(f"category_id {category_id} 的数量为: {count}")

    return category_id_counts

# 定义文件夹路径
folder_path = '/codeyard/yolov5_6.0/data_MS_ALLlabels'

# 如果要统计特定的 category_id，将其作为参数传递
target_category_id = None

# 如果 target_id 不为 None，则统计特定的 category_id，否则统计所有 category_id
category_counts = count_category_ids(folder_path, target_id=target_category_id)

上述程序定义了一个名为count_category_ids的函数，该函数接受一个文件夹路径作为输入，并可选择性地传递一个目标类别ID进行统计。

程序首先获取指定文件夹中所有以coco_filtered.json结尾的JSON文件，并将其存储在json_files列表中。然后，它创建一个空字典category_id_counts，用于存储各个类别ID的数量。

接下来，程序循环遍历每个JSON文件，读取文件内容，并开始统计各个类别ID的数量。对于每个注释（annotation）对象，它会提取category_id字段的值，并根据目标类别ID的选择进行统计。如果目标类别ID为None，则统计所有类别ID；否则，只统计与目标类别ID相等的类别ID。

每次出现一个新的类别ID时，程序将在category_id_counts字典中创建对应的键，并将值初始化为0。然后，它会增加相应类别ID的计数器。

最后，程序会打印每个类别ID及其对应的数量。