分层采样划分数据集

import glob

cats_fenbu = {}
for i in range(18):
    cats_fenbu[str(i)] = 0

# 统计每个类别的个数
for txt in glob.glob(r"D:\dataset\TianChi_clothes\guangdong1_round1_train2_20190828\train\labels\*"):
    with open(txt, "r") as f:
        labels = f.read().splitlines()
        for line in labels:
            cat = line.split(" ")[0]
            cats_fenbu[str(cat)] += 1
print(cats_fenbu)

# 按8:2划分训练测试数据集,直接产生测试集，剩余的就是训练集
cats_test = {}      # 存放测试集每个类别应有的个数
for cats_fenbu_key in cats_fenbu:
    cats_test[cats_fenbu_key] = round(cats_fenbu[cats_fenbu_key] / 10 * 2)
print(cats_test)

cats_test_real = {}     # 存放测试集中当前已加入的各类别数
for i in range(18):
    cats_test_real[str(i)] = 0
dataset_test = []          # 存放测试集的图片名
for txt in glob.glob(r"D:\dataset\TianChi_clothes\guangdong1_round1_train2_20190828\train\labels\*"):
    with open(txt, "r") as f:
        labels = f.read().splitlines()
        txt