如何基于自定义不平衡图像数据集构建等比例分层抽样测试集？

核心思路 #

要得到4类占比完全相同的2万张测试集，本质是按类别分层抽样，每个类别抽取5000张（20000 ÷ 4），基于你已有的CSV文件可以快速实现。

具体步骤与代码 #

1. 加载并检查原始数据
   先读取CSV文件，确认每个类别的样本量是否满足抽取需求（每个类别至少5000张）：

import pandas as pd

# 替换为你的CSV文件路径
df = pd.read_csv('your_image_metadata.csv')

# 查看各类别样本数量
class_dist = df['类别'].value_counts()
print("原始类别分布：")
print(class_dist)

若某类别样本数不足5000，需调整抽取策略（比如缩减总测试集规模，或和需求方确认是否允许该类用全部样本）。

2. 按类别分层抽样
   对每个类别单独抽取指定数量的样本，保证类别平衡：

# 每个类别抽取5000张
sample_size_per_class = 20000 // 4

# 分组抽样，random_state保证结果可复现
balanced_test_df = df.groupby('类别').apply(
    lambda group: group.sample(n=sample_size_per_class, random_state=42)
)

# 重置索引，去除groupby带来的层级索引
balanced_test_df = balanced_test_df.reset_index(drop=True)

3. 验证并保存结果
   检查抽样后的类别分布，确认符合要求后保存为新的CSV：

# 验证抽样后的类别分布
print("\n抽样后类别分布：")
print(balanced_test_df['类别'].value_counts())

# 保存平衡后的测试集CSV，index=False避免生成额外索引列
balanced_test_df.to_csv('balanced_test_set.csv', index=False)

注意事项 #

• 可复现性：设置random_state固定抽样结果，方便后续重复实验或调整。
• 路径有效性：抽样完成后，可随机抽查几条记录的路径字段，确认能正常访问对应图像文件。
• 特殊情况处理：若某类别样本量不足，可改为抽取该类全部样本，同时按比例减少其他类的抽取数量（比如某类只有3000张，总测试集改为12000张，每类3000张）。

内容的提问来源于stack exchange，提问作者hailey