cnn-captcha/tools/collect_labels.py

#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
统计样本的标签，并写入文件labels.json
"""
import os
import json
from log_ware import LogWare

logger = LogWare().get_logger()

image_dir = "../sample/origin"
image_list = os.listdir(image_dir)

labels = set()
for img in image_list:
    split_result = img.split("_")
    if len(split_result) == 2:
        label, name = split_result
        if label:
            for word in label:
                labels.add(word)
    else:
        pass

logger.debug("共有标签%d种", len(labels))

with open("./labels.json", "w") as f:
    f.write(json.dumps("".join(list(labels)), ensure_ascii=False))

logger.debug("将标签列表写入文件labels.json成功")