# 统计标注好的数据，同时给出统计结果保存为txt
import os
def count_imgs(path:str, tag:str)->str:
    target_path = os.path.join(path, tag)
    # 获取类别子目录
    sta_res = {}
    for c in os.listdir(target_path):
        cls_path = os.path.join(target_path, c)
        # 获取图片
        imgs = os.listdir(cls_path)
        sta_res[c] = len(imgs)
    return f'{tag} data statistics: {sta_res}'
def main():
    train_data_path = r'D:\code\water_turbidity_det\label_data'
    dirs = os.listdir(train_data_path)
    info = []
    if 'train' in dirs:
        info.append(count_imgs(train_data_path, 'train'))
    if 'test' in dirs:
        info.append(count_imgs(train_data_path, 'test'))
    if 'val' in dirs:
        info.append(count_imgs(train_data_path, 'val'))
    with open(os.path.join(train_data_path, 'statistic.txt'), 'w') as fw:
        for i in info:
            fw.write(i)
            fw.write('\n')
if __name__ == '__main__':
    main()