YOLO数据集标签数量统计脚本

在目标检测模型的开发流程中，数据集的质量直接决定了模型的上限。而类别分布作为数据集质量的核心指标之一，更是我们必须重点关注的内容 —— 如果某些类别的样本数量极少，模型可能难以学到其特征；如果类别分布极度不平衡，训练出的模型会倾向于预测占比高的类别。今天，我们就来详细解析一个实用的 YOLO 格式数据集标签统计脚本，看看它是如何帮我们快速掌握类别分布的。

一、为什么需要统计标签数量？

在开始解析代码前，我们先明确一个问题：为什么要花时间做标签数量统计？

对于 YOLO（或其他目标检测框架）的数据集来说，标签文件（通常是.txt格式）记录了每个目标的类别和位置信息。其中，类别索引是标签文件的第一列数据，直接决定了这个目标属于哪个类别。统计每个类别出现的总次数，能帮我们：

1. 发现类别不平衡问题：比如 “猫” 有 10000 个样本，而 “狗” 只有 100 个，模型训练时很可能 “偏爱” 预测 “猫”；

2. 评估数据集覆盖度：检查是否有遗漏的类别（比如标注规范里有 10 类，但实际只标注了 8 类）；

3. 指导数据增强策略：对样本少的类别，我们可以通过旋转、裁剪等增强方式扩充数据；

4. 验证标注一致性：如果某个类别在不同子集中（如 train/val）的占比差异过大，可能是标注错误。

因此，一个可靠的标签统计脚本，是数据集分析的 “第一块拼图”。

二、脚本整体概览

我们今天解析的脚本名为labels_count.py，核心功能是遍历 YOLO 格式数据集的标签文件夹，统计每个类别出现的总次数，并输出详细结果。

先来看脚本的整体代码：

import os
def labels_count(root):'''统计数据集中每个类别的数量args:root:yolo格式数据集的label路径return:labels_count:数据集中每个类别的数量'''labels_count = {}for dir in os.listdir(root):labels_dir_count = {}if dir in ['labels']:label_dir = os.path.join(root, dir)label_files = os.listdir(label_dir)  # 展示目标文件夹下所有的文件名label_files = list(filter(lambda x: x.endswith('.txt'), label_files))  # 取到所有以.txt结尾的yolo格式文件for label_file in label_files:with open(os.path.join(label_dir, label_file), 'r') as f:lines = f.readlines()for line in lines:line = line.strip().split()class_name = line[0]if class_name in labels_count:labels_count[class_name] += 1else:labels_count[class_name] = 1if class_name in labels_dir_count:labels_dir_count[class_name] += 1else:labels_dir_count[class_name] = 1print(f'{dir} done!', 'labels_count:', labels_dir_count)return labels_countif __name__ == '__main__':root = r'data/labels' # dataset为你的YOLOv8格式的数据集labels_cnt = labels_count(root)print('all sets done! labels_count:', labels_cnt)

整体逻辑非常清晰：通过os模块处理文件路径和读写，用labels_count函数实现核心统计逻辑，最后在主函数中指定数据集路径并执行统计。

三、核心代码逐行解析

接下来，我们一步步拆解labels_count函数的实现细节，搞懂它是如何完成统计工作的。

1. 函数定义与参数说明

def labels_count(root):'''统计数据集中每个类别的数量args:root:yolo格式数据集的label路径return:labels_count:数据集中每个类别的数量'''

函数labels_count接收一个参数root，表示 YOLO 数据集标签文件夹的根路径（比如data/labels），返回一个字典labels_count，键是类别索引（字符串格式），值是该类别的总出现次数。

这里的文档字符串（docstring）很重要，它清晰说明了函数的作用、参数和返回值，符合 Python 代码的规范，方便其他开发者理解和使用。

2. 初始化计数字典

labels_count = {}  # 存储所有标签文件夹的总计数

我们用一个空字典labels_count来记录所有类别在整个数据集中的总出现次数。字典的键是类别索引（比如'0'、'1'，对应 YOLO 标签中第一列的字符串），值是累计次数。

3. 遍历根目录，筛选目标文件夹

for dir in os.listdir(root):labels_dir_count = {}  # 存储当前标签文件夹的计数if dir in ['labels']:  # 只处理名为'labels'的文件夹label_dir = os.path.join(root, dir)  # 拼接完整的标签文件夹路径

• os.listdir(root)：列出root路径下的所有文件和文件夹（如['labels', 'images', 'classes.txt']）；

• if dir in ['labels']：YOLO 数据集的标签文件通常集中在labels文件夹中，因此这里只处理名为labels的文件夹（如果你的数据集分train/labels、val/labels，可以修改这里的列表，比如['train/labels', 'val/labels']）；

• os.path.join(root, dir)：拼接路径（比如root是data，dir是labels，则结果为data/labels），避免直接用字符串拼接（如root + '/' + dir）导致的跨平台问题（Windows 用\，Linux 用/）。

4. 筛选标签文件（.txt 格式）

label_files = os.listdir(label_dir)  # 列出标签文件夹下的所有文件
label_files = list(filter(lambda x: x.endswith('.txt'), label_files))  # 筛选出.txt文件

• YOLO 的标签文件以.txt为后缀（每个图片对应一个同名的.txt标签文件），因此需要用filter函数筛选出所有.txt文件；

• lambda x: x.endswith('.txt')是一个匿名函数，用于判断文件名是否以.txt结尾，配合filter实现筛选。

5. 读取标签文件，统计类别次数

这部分是脚本的核心，负责逐行解析标签文件，提取类别索引并更新计数：

for label_file in label_files:  # 遍历每个标签文件# 打开标签文件并读取所有行with open(os.path.join(label_dir, label_file), 'r') as f:lines = f.readlines()# 逐行解析标签内容for line in lines:line = line.strip().split()  # 去除空白字符，按空格分割class_name = line[0]  # 第一列是类别索引# 更新总计数字典if class_name in labels_count:labels_count[class_name] += 1else:labels_count[class_name] = 1# 更新当前文件夹的计数字典if class_name in labels_dir_count:labels_dir_count[class_name] += 1else:labels_dir_count[class_name] = 1

这里有几个关键细节需要注意：

• 标签文件格式：YOLO 的标签文件中，每一行代表一个目标，格式为class_id x_center y_center width height（如0 0.5 0.5 0.8 0.6，表示类别 0 的目标，中心坐标和宽高）。因此line[0]就是类别索引；

• strip()和split()：strip()用于去除行首行尾的空白字符（如换行符、空格），split()按空格分割字符串，得到一个列表（如上述例子会分割为['0', '0.5', '0.5', '0.8', '0.6']）；

• 双字典计数：labels_count统计所有文件夹的总次数，labels_dir_count统计当前labels文件夹的次数，方便我们既看全局又看局部。

6. 打印当前文件夹的统计结果

print(f'{dir} done!', 'labels_count:', labels_dir_count)

当一个labels文件夹处理完成后，打印该文件夹的类别计数（如labels done! labels_count: {'0': 100, '1': 200}），方便我们实时查看进度和中间结果。

7. 主函数：执行统计并输出总结果

if __name__ == '__main__':root = r'data/labels'  # 数据集标签路径（可修改为自己的路径）labels_cnt = labels_count(root)print('all sets done! labels_count:', labels_cnt)

主函数的作用是指定数据集路径（root），调用labels_count函数，并打印最终的总统计结果。这里的r'data/labels'是原始脚本的示例路径，实际使用时需要修改为自己的数据集标签文件夹路径（比如r'E:/datasets/coco/labels'）。

四、如何使用这个脚本？

使用步骤非常简单，只需 3 步：

1. 准备数据集：确保你的数据集是 YOLO 格式，标签文件（.txt）存放在labels文件夹中；

2. 修改路径：在主函数中，将root = r'data/labels'改为你的标签文件夹实际路径（注意 Windows 路径用\或/，或在字符串前加r避免转义问题）；

3. 运行脚本：执行python labels_count.py，控制台会输出每个文件夹的统计结果和总结果。

例如，假设你的数据集有 3 个类别，运行后可能得到这样的输出：

labels done! labels_count: {'0': 500, '1': 300, '2': 200}
all sets done! labels_count: {'0': 500, '1': 300, '2': 200}

这说明labels文件夹中，类别 0 出现 500 次，类别 1 出现 300 次，类别 2 出现 200 次。

五、进阶扩展：让脚本更实用

原始脚本已经能满足基础的统计需求，但在实际场景中，我们可以根据需要扩展功能，让它更强大。

1. 支持多子集统计（train/val/test）

实际数据集通常会分为训练集（train）、验证集（val）、测试集（test），每个子集都有自己的labels文件夹（如train/labels、val/labels）。我们可以修改代码，同时统计多个子集：

# 修改筛选条件，支持多个标签文件夹
if dir in ['train/labels', 'val/labels', 'test/labels']:# 后续逻辑不变

这样就能分别得到每个子集的类别分布，方便对比训练 / 验证集的一致性。

2. 映射类别索引到类别名称

YOLO 数据集通常会有一个classes.txt文件，记录类别索引对应的名称（如0: person、1: car）。我们可以读取这个文件，将统计结果中的索引转换为名称，更直观：

# 读取classes.txt，建立索引到名称的映射
def get_class_names(classes_path):with open(classes_path, 'r') as f:return [line.strip() for line in f.readlines()]
​
# 在labels_count函数中使用映射
class_names = get_class_names(os.path.join(root, 'classes.txt'))
# 统计时用class_names[int(class_name)]作为键

修改后，统计结果会变成{'person': 500, 'car': 300}，更易读。