YOLO系列之统计各类型标签数量

YOLO系列之统计各类型标签数量

一句话总结：本文用于统计 YOLOv8格式数据集中各类型 标签数量。

前言

本系列的其他博客——YOLO系列之数据集划分（以YOLOv8为例），有需要可以参考。

YOLO系列数据集格式

在数据集中，图像文件存放于images目录下，对应的YOLO格式的txt标签文件要存放于labels目录下，而且要在images和labels目录下分别建立train训练集、valid验证集、test
测试集（可选，但强烈建议要有！）子目录，如下展示：

dataset
	- images
		- train
			- 1.jpg
			- 5.jpg # 因为要随机划分，所以数字是随机写的，下面也一样
			-  ... 
		- valid
			- 2.jpg
			- 7.jpg
			-  ...
		- test
			- 3.jpg
			- 9.jpg
			- ...
	- labels
		- train
			- 1.txt
			- 5.txt
			-  ... 
		- valid
			- 2.txt
			- 7.txt
			-  ...
		- test
			- 3.txt
			- 9.txt
			- ...

实际目录如下图：

代码实现

以下代码可以直接复制运行（只需要修改root路径）

本代码会输出训练集train、验证集valid和测试集test中各类型标签数量，格式为数据集（如train） done! labels_count: {'标签编号如（0）': 标签数量}，如果是多目标的话，则为数据集（如train） done! labels_count: {'0': 111， ‘1’:222, '2':333}。
最后输出总的数量：all sets done! labels_count: {'标签编号如（0）': 标签数量}。
代码如下：

import os
def labels_count(root):
    '''
    统计数据集中每个类别的数量
    args:
        root:yolog格式数据集的label路径
    return:
        labels_count:数据集中每个类别的数量
    '''
    labels_count = {}
    for dir in os.listdir(root):
        labels_dir_count = {}
        if dir in ['train', 'valid', 'test']:
            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] += 1
                    else:
                        labels_count[class_name] = 1
                    if class_name in labels_dir_count:
                        labels_dir_count[class_name] += 1
                    else:
                        labels_dir_count[class_name] = 1
            print(f'{dir} done!', 'labels_count:', labels_dir_count)

    return labels_count

if __name__ == '__main__':
    root = r'dataset\labels' # dataset为你的YOLOv8格式的数据集
    labels_cnt = labels_count(root)
    print('all sets done! labels_count:', labels_cnt)

在我的数据集上输出如下：

test done! labels_count: {'0': 874}
train done! labels_count: {'0': 7027}
valid done! labels_count: {'0': 869}
all sets done! labels_count: {'0': 8770}