Python统计标签个数

深度学习中的图像处理，包含多个标签的数据集，往往存在类别不平衡的问题，这样会导致训练模型准确率低的问题，因此有必要统计，各个类别图片标签个数，以及所占的比例。
举例：现有标签文件train.txt，内容如下：
/data/train/1.jpg 4
/data/train/88.jpg 2
/data/train/455.jpg 1
/data/train/764.jpg 3
/data/train/64.jpg 4
…

code：

def count_label(data):
    total = 0
    ratio = 0.0
    data_dict = {}

    file = open(data,'r')
    lines = file.readlines()

    for line in lines:
        content = line.strip().split()
        label = content[1]

        # 将标签和标签所对应的图片数目统计为dict格式:{label：num_label}
        if label not in data_dict:
            data_dict[label] = 0
        data_dict[label] = data_dict[label] + 1

    # 统计标签的总数目
    for val in data_dict.values():
        total += val

    # 计算每个标签所占的比例ratio,并将原来dict中的num_label替换为ratio,即{label:ratio}
    for key,val in data_dict.items():
        ratio = float(val/total)
        data_dict[key] = ratio
        
    return data_dict

count_label_dict = count_label('train.txt')
print(count_label_dict)